前言
最近一直在研究代码审计agent相关的东西,试图用pi和pi-durable去实现一些更好的效果,有点好奇pi去读代码的时候到底在干什么,于是就有了这个
Pi在干什么
你知道的,pi一向以极简著称,在宣传会上开发者表面只有read write bash edit,可以看看pi的工具集
export const allToolNames: Set<ToolName> = new Set([ "read", "bash", "powershell", "edit", "write", "grep", "find", "ls",]);pi的核心是一个简单的react循环,agent-loop.ts里面写的比较清楚
/** * Main loop logic shared by agentLoop and agentLoopContinue. */async function runLoop( initialContext: AgentContext, newMessages: AgentMessage[], initialConfig: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, streamFunction: StreamFn,): Promise<void> { let currentContext = initialContext; let config = initialConfig; let lastCompletedTurn: PrepareNextTurnContext | undefined; let explicitContinuation = false; // Check for steering messages at start (user may have typed while waiting) let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
// Outer loop: continues when queued follow-up messages arrive after agent would stop while (true) { let hasMoreToolCalls = true;
// Inner loop: process tool calls and steering messages while (hasMoreToolCalls || pendingMessages.length > 0) { let preparedMessages: AgentMessage[] = []; if (lastCompletedTurn) { const nextTurnSnapshot = await config.prepareNextTurn?.(lastCompletedTurn); if (nextTurnSnapshot) { currentContext = nextTurnSnapshot.context ?? currentContext; preparedMessages = nextTurnSnapshot.messages ?? []; config = { ...config, model: nextTurnSnapshot.model ?? config.model, reasoning: nextTurnSnapshot.thinkingLevel === undefined ? config.reasoning : nextTurnSnapshot.thinkingLevel === "off" ? undefined : nextTurnSnapshot.thinkingLevel, }; } // Preparation can be long-running (for example, compaction). Pick up steering // queued while it ran. Only poll again if the earlier poll returned nothing; // otherwise one-at-a-time mode would deliver two messages in this turn. if (pendingMessages.length === 0) { pendingMessages = (await config.getSteeringMessages?.()) || []; } await emit({ type: "turn_start" }); }
// Process prepared and queued messages before the next assistant response. for (const message of declareToolChanges(currentContext, [...preparedMessages, ...pendingMessages])) { await emit({ type: "message_start", message }); await emit({ type: "message_end", message }); currentContext.messages.push(message); newMessages.push(message); } pendingMessages = [];
const requestUpdate = await config.prepareRequest?.( { context: currentContext, model: config.model, thinkingLevel: config.reasoning ?? "off", }, signal, ); if (requestUpdate) { currentContext = requestUpdate.context ?? currentContext; config = { ...config, model: requestUpdate.model ?? config.model, reasoning: requestUpdate.thinkingLevel === undefined ? config.reasoning : requestUpdate.thinkingLevel === "off" ? undefined : requestUpdate.thinkingLevel, }; }
// Stream assistant response const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFunction); newMessages.push(message);
if (message.stopReason === "error" || message.stopReason === "aborted") { lastCompletedTurn = { message, toolResults: [], context: currentContext, newMessages, }; await config.finishTurn?.(lastCompletedTurn, signal); await emit({ type: "turn_end", message, toolResults: [] }); await emit({ type: "agent_end", messages: newMessages }); return; }
// Check for tool calls const toolCalls = message.content.filter((c) => c.type === "toolCall");
const toolResults: ToolResultMessage[] = []; hasMoreToolCalls = false; if (toolCalls.length > 0) { // A "length" stop means the output was cut off by the token limit, so // every tool call in the message may carry truncated arguments. Fail // them all instead of executing potentially borked calls. const executedToolBatch = message.stopReason === "length" ? await failToolCallsFromTruncatedMessage(toolCalls, emit) : await executeToolCalls(currentContext, message, config, signal, emit); toolResults.push(...executedToolBatch.messages); hasMoreToolCalls = !executedToolBatch.terminate;
for (const result of toolResults) { currentContext.messages.push(result); newMessages.push(result); } }
lastCompletedTurn = { message, toolResults, context: currentContext, newMessages, }; const decision = await config.finishTurn?.(lastCompletedTurn, signal); await emit({ type: "turn_end", message, toolResults });
if (decision?.action === "end") { await emit({ type: "agent_end", messages: newMessages }); return; }
explicitContinuation = decision?.action === "continue"; pendingMessages = (await config.getSteeringMessages?.()) || []; if (hasMoreToolCalls || pendingMessages.length > 0) { explicitContinuation = false; } }
// Agent would stop here. Check for follow-up messages. const followUpMessages = (await config.getFollowUpMessages?.()) || []; if (followUpMessages.length > 0) { // Set as pending so inner loop processes them explicitContinuation = false; pendingMessages = followUpMessages; continue; }
// No natural request was selected, so fulfill the continuation decision with one context-only turn. if (explicitContinuation) { explicitContinuation = false; continue; }
// No more messages, exit break; }
await emit({ type: "agent_end", messages: newMessages });}他会不断的循环,每次循环都类似一次LLM调用+tools调用,返回结果ToolResultMessage append 进 transcript,从而在下一轮被看见
同时可以看见停止要求是
hasMoreToolCalls = !executedToolBatch.terminate;
通过这个去判断是否还有活要干,如果没有活了就停了,下面看看让他去读一个项目是怎么跑的
首先会加载提示词
preamble "You are an expert coding assistant operating inside pi..." tools - read: Read file contents - bash: Execute shell commands ... rules - Use read to examine files instead of cat or sed. ← 工具自己贡献的 - Be concise in your responses project_context <project_instructions path="AGENTS.md">...</...> skills ... cwdagent拿到提示词后会转换成user的输入agentLoop(prompts, context, config, signal, streamFn)
在read.ts中会限制每次读的大小DEFAULT_MAX_LINES=2000 / DEFAULT_MAX_BYTES=50KB,然后就按照read工具所写的一样,去进行分片阅读,没有自动索引、没有 AST、没有依赖图
那让他去代码审计漏洞挖掘的时候,其实也没有去理解这个项目,而是去匹配一些危险函数比如eval什么的,然后到那个代码的附近去read,然后再跳转到引用到的其他的函数,然后不断的重复,这就是一个裸pi的操作
这样明显是不够的,如果代码审计没有理解项目,一些漏洞是一辈子审计不出来的,这个时候就需要弥补一下他的缺点
Pi的记忆模块
pi的记忆结构是一个tree,可以看看我的另一个帖子https://www.zhuangsanmeng.xyz/posts/pi-tree/,补充一下来说fork一个会话=从某个历史entry长出新leaf。
在session-manager.ts中写了记忆的读取方法
const contextEntries = [compaction]; // 摘要进 for (let i = 0; i < compactionIdx; i++) { if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; if (foundFirstKept && entry.message.role !== "system") contextEntries.push(entry); } contextEntries.push(...path.slice(compactionIdx + 1));整个简单来说就是
entries[] ──buildSessionPath()──→ path[] 按 leaf 回溯出的线性路径 ──buildContextEntries()→ contextEntries 折叠 compaction ──应用 context_edit──→ messages[] 最终喂给 LLM比较好玩的是tui里面会保存全部的记忆,但是llm只能看到最后一部分,见compaction.ts
DEFAULT_COMPACTION_SETTINGS = { enabled: true, reserveTokens: 16384, // 预留 keepRecentTokens: 20000 // 保留最近 }实际触发在prepareNextTurn → _compactBeforeNextAssistantResponse(),也就是两次 LLM 调用之间。切点由findProjectedCutPoint()从后往前累加 token,攒够 20000 就切。
还有跨会话记忆模块可以基于subagent去进行管理,这个是有助于渗透和代码审计的(
一些优化
首先要解决如何让pi去理解项目,这里的方法是让subagent主动搭建docker完整环境,然后登录、注册等行为去记录操作和这个系统的用处,然后共享给主agent和其他的subagent
至于read的解决方法,这里选择尽可能的去分片完整阅读,记录未阅读的代码区块,并且也要去识别危险函数等,期间着重阅读主要相关的代码。
契约确认 → navigator → 台账播种 ↓ 分波区块审计(audit-authz / audit-input / audit-static,按关注点分配) ↓ 证据门(回源核对引用 + 主动找漏掉的守卫 + 严判 needs_runtime) ↓ 交叉审计(只针对 medium+ / 跨信任边界 / 低置信度) ↓ 终态:confirmed | rejected | uncertain | duplicate ↓ 双门禁校验 → 报告一个比较简单的思路
关于agent自建
其实我一直在想,现在写agent,我们本质是在写什么,是harness?还是agent的调度器
我选择了后者
在这样一个pi、codex、claudecode…很强的时代,我认为手写一个harness是没有必要的,我们只需要去把一个现有的agent调优成针对性强的agent就行了
Pi-durable
pi开发团队的一个新项目,应该说是从pi中抽象出来的一个层面?
更像是为了agent开发准备的,他有模型请求、工具调用、压缩等任务的checkpoint恢复,有steering指令,有自建subagent管理方法等等,反正一个agent需要的东西都有,最近在研究这个去写一个agent的核心调度层,虽然后期api可能会变,但是感觉还是值得学习的,可以详细看看原文,就不重复了(
目前看官方的发文,是非常适合长任务的,let’s go
近期开发
最近其实在开发中我在研究一些乱七八糟的东西,比如代码审计需要什么样的辅助,人工审计的时候,可以打端点、lsp跳转,这些对于ai来说是不太现实的,于是我在尝试用类似openwiki的方法,把一个需要挖掘的项目抽象成一个代码图谱和wiki,然后subagent去分链路交叉审计,哎,这样是不是有点像ARTEX的双图结构呢?好像是有点像的,hh
问题其实也很明显,代码审计并不是渗透一样的信息收集过程,白盒代码审计更重要的是不同链子之间的交叉,比如xss组合拳等,ai在一些方面不太敏感,或者是在长任务中不断的压缩就会丢失,于是我加入了每个subagent的时间限制、vcc算法压缩、obverse的recall机制,但是这样貌似太重了,有没有什么更加轻量化的方法呢,还在不断地探索。。
未完待续吧()
部分信息可能已经过时