diff --git a/AGENTS.md b/AGENTS.md index 978c9d2..3705d9f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,12 +87,12 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del - **Channels** publish inbound messages through `MessageBus`; outbound writes arrive through `OutboundDispatcher` or a per-turn `TurnSink`. They know nothing about sessions or LLM - **Inbound contract** carries normalized sender/time/media plus `ChannelContext`; core routing may interpret `reply_to` but must treat platform-private context as opaque reply data - **MessageBus** owns bounded inbound/outbound/control queues; outbound routing and retries belong to `OutboundDispatcher`, not the queue -- **SessionManager** owns session state, dialog operations, context construction, per-session work queues, and persistence coordination +- **SessionManager** owns session state, dialog operations, context construction, per-session work queues, active-Turn steering admission, and persistence coordination - **TurnController** is the only owner of active Turn state; snapshots are complete latest-wins values, not token queues, and `Completed` is published only after atomic persistence succeeds - **DeliveryCoordinator** projects active Turn snapshots without mutating history; it owns `TurnSink` lifecycle but no platform message IDs, which remain private to each sink - **WorkManager** owns the single active plan per session, item state transitions, plan versions, and plan-change events; plans are optional and absent from ordinary chat context - **Scheduler** supports legacy direct-delivery jobs and managed `task`/`monitor` jobs; managed agents cannot call `send_message`, and `on_alert` suppresses only healthy informational results -- **AgentLoop** is stateless across turns; it receives prepared history, calls LLM providers, executes tools, and returns one result +- **AgentLoop** is stateless across turns; it receives prepared history, drains same-Turn steering only at safe model boundaries, calls LLM providers, executes tools, and returns one result - **WebUI management APIs** only expose allowlisted config/profile/log/storage operations; keep response limits, secret redaction, atomic config writes, and profile path allowlists intact - **WebUI styling** uses the local Fluent 2 semantic tokens in `webui/src/styles.css`; page and component styles should consume the aliases instead of introducing independent hard-coded palettes, and must preserve selectable light/dark and brand-color themes, keyboard focus, responsive layout, and reduced-motion behavior. Browser-only appearance settings belong in `lib/theme.js`/`localStorage`, and `public/theme-init.js` must restore them before Svelte mounts - **Gateway config reload** validates and prepares a complete next runtime generation before retiring the old one; close runtime admission before draining inbound lanes, Sessions, Scheduler jobs, and background sub-agents; MCP connects only during activation; host/port, workspace, and effective storage path remain restart-only, and runtime reload must never mutate process environment variables @@ -111,7 +111,8 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del ### Concurrency and Lifecycle Invariants -- Messages in one session are processed serially through a bounded queue; different sessions may run concurrently +- One session runs at most one Turn; ordinary input steers its active Turn by default, `/queue` explicitly waits for the next Turn, and different sessions may run concurrently +- Steering admission, final close, fallback, and `/stop` must be lossless and mutually exclusive: an input belongs to exactly the active Turn or the next-Turn FIFO, while `/stop` intentionally discards both - Outbound messages are ordered per `(channel, chat_id)`; a slow destination must not block unrelated destinations - Active Turn delivery and ordinary outbound delivery share the same per-`(channel, chat_id)` write lock; never enqueue token deltas into MessageBus - Slow Turn consumers may skip intermediate snapshots but must receive an explicit bounded terminal delivery; shutdown must call sink abort so platform cleanup remains possible diff --git a/Cargo.toml b/Cargo.toml index 1790b7a..7251924 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "picobot" -version = "1.4.1" +version = "1.5.0" edition = "2024" [dependencies] diff --git a/README.md b/README.md index 329be25..aa10b00 100644 --- a/README.md +++ b/README.md @@ -256,7 +256,7 @@ WebUI/TUI 上传文件默认保存到 `~/.picobot/media/cli_chat`，单文件上 详细时序和失败语义见 [架构文档：消息与控制数据流](docs/ARCHITECTURE.md#4-消息与控制数据流)。 -同一 session 的普通消息由专属有界队列串行处理，不同 session 可以并发；活动 Turn 使用 latest-wins 快照，慢展示端只跳过中间状态，不反压模型。普通出站消息按 `(channel, chat_id)` 分 lane 保序，两条投递路径共享目标写锁。Gateway 的长生命周期任务统一由 `TaskSupervisor` 取消和限时回收。 +同一 session 始终只运行一个 Turn，不同 session 可以并发。Turn 执行期间新发的普通消息默认 steering 当前工作：系统在完整工具批次后或最终回复边界把它作为真实用户消息加入下一次模型调用；使用 `/queue ` 可明确等当前 Turn 完成后再处理，使用 `/stop` 可中断当前 Turn 并清空等待输入。Steering mailbox 和 session 队列都有界且带可靠回退。活动 Turn 使用 latest-wins 快照，慢展示端只跳过中间状态，不反压模型。普通出站消息按 `(channel, chat_id)` 分 lane 保序，两条投递路径共享目标写锁。 核心边界： @@ -311,6 +311,7 @@ Session ID 使用三段式： | `/dump` | 导出当前 dialog 为 Markdown | | `/mcp` | 查看 MCP 服务器和工具状态 | | `/health` | 检查 PicoBot 运行依赖 | +| `/queue ` | 等当前 Turn 完成后再把消息作为下一 Turn 处理 | | `/stop` | 停止当前任务并清空队列 | | `/todo [done\|cancel]` | 查看、完成或取消当前 session 的任务计划 | | `/reload` | 校验并重新加载 Gateway 配置 | diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 3a0f1e7..b88db32 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -107,11 +107,16 @@ sequenceDiagram C->>B: publish InboundMessage B->>G: consume inbound G->>S: handle_message - S->>W: try_send AgentTask + S->>W: try_send AgentTask (idle or /queue) S-->>G: AgentProcessing W->>T: start Turn W->>L: subscribe latest snapshots W->>A: process_streaming(history) + C->>B: ordinary input during active Turn + B->>G: consume inbound + G->>S: handle_message + S->>A: bounded steering mailbox + A->>A: drain after tool batch / before final A-->>T: reasoning/text/tool events T-->>L: complete TurnSnapshot L->>C: TurnSink update (best effort) @@ -127,11 +132,14 @@ sequenceDiagram 关键语义： -- Gateway inbound router 按 `(channel, chat_id)` 使用容量 32 的短生命周期 lane 保持入口顺序，不同聊天可并发路由；普通消息进入对应 session worker 后立即返回 `AgentProcessing`。 +- Gateway inbound router 按 `(channel, chat_id)` 使用容量 32 的短生命周期 lane 保持入口顺序，不同聊天可并发路由；没有活动 Turn 时，普通消息进入对应 session worker 后立即返回 `AgentProcessing`。 - `/stop` 绕过同聊天的 inbound lane，直接使正在运行的 worker/Turn 失效；其他 slash command 仍在聊天 lane 内有序执行。 -- 每个 session 有一条容量为 32 的队列，同一 session 串行处理，不同 session 的 worker 可并发执行。 +- 活动 Turn 存在时，普通消息默认作为 steering 进入本 Turn 的有界 mailbox；`/queue ` 明确进入下一 Turn。AgentLoop 只在完整工具批次结束后、或准备接受无工具最终回复时排空 mailbox，并把输入作为真实、可持久化的 `role=user` 消息加入下一次模型请求。 +- Steering mailbox 最多容纳 32 条、合计 64 KiB 文本与元数据。mailbox 已关闭或满时，输入可靠回退到 session 队列；若 session 队列也满则明确拒绝。Session 在入站时分配单调序号，Turn 结束时未消费的 steering 由 worker 本地恢复队列接管，并与 `/queue` 输入按该序号合并选择，不能丢失或互相超越。 +- 每个 session 有一条容量为 32 的普通队列，同一 session 仍只运行一个 Turn，不同 session 的 worker 可并发执行。 - 队列满时明确拒绝新消息，不允许无界积压。 -- Slash command 不进入 Agent 队列，由 `SessionManager` 直接执行，因此 `/stop` 等控制操作不会排在长模型调用之后。 +- Slash command 通常不进入 Agent 队列，由 `SessionManager` 直接执行；`/queue` 是显式排队输入，`/stop` 是显式中断并清空当前 mailbox 与队列。 +- WebSocket `user_input.client_message_id` 只用于让 `turn_committed` 以同一消息 ID 替换 WebUI 的乐观用户气泡；它不改变入站顺序或 steering/queue 决策。 - `InboundMessage` 只保存规范化输入：`sender_id`、`received_at`、媒体和一个 `ChannelContext`。核心只解释其中的 `reply_to`，其语义是本轮出站应回复的当前入站消息；被用户引用的父消息只用于补充模型上下文。reaction/message ID、话题 root/thread 等平台字段作为 `private` 不透明传到对应 Turn/普通回复，不能散落为核心层 magic key。持久化的用户消息保留真实接收时间和 `UserInput` 来源，客户端历史投影不暴露来源中的平台用户 ID。 - Session 为每个 Agent 请求创建一个 `TurnController`。Provider 向 AgentLoop 发 delta，AgentLoop 发结构化 TurnEvent，只有 TurnController 能把事件归约为有序 block 和单调 revision 的完整快照。 - Turn 快照经 Tokio `watch` 发布，语义为 latest-wins；慢客户端或慢渠道跳过中间状态，不反压 Provider。终态明确编码在快照中，不依赖 sender 关闭。 @@ -192,10 +200,11 @@ Session ID 格式为： 4. 慢操作开始前记录 `state_version`，提交前重新验证，防止旧快照覆盖 `/clear`、`/delete` 等并发修改。 5. 持久化写入由 `persistence_lock` 串行化；多条相关记录应使用 Storage 的原子接口。 6. 内存先变更但持久化失败时，必须回滚精确匹配的消息后缀，不能删除无关的新状态。 +7. Steering 的接收、最终边界关闭和 `/stop` 必须通过同一个 mailbox 状态串行化；每条输入只能落入当前 Turn 或下一 Turn 之一。 当前 WebUI/TUI Turn 通过 `send_message(files=...)` 向自身 session 投递文件时，文件先进入 task-local Turn delivery 暂存区，成功结束后附加到最终 assistant 消息，与工具链一起原子提交；因此持久化和刷新后的顺序都是工具调用/结果在前、携带附件的最终回复在后，也不会生成带 `[message from ...]` 的自投递气泡。其他同 Turn 自投递仍是受控例外：只有 task-local Turn ID 仍匹配该 session 的 active Turn，写入才允许不递增 `state_version`。跨 Turn、跨 session 以及无法证明所有权的写入仍必须递增版本。Provider 回放历史附件时，只有 user 输入和当前工具结果可生成模型原生媒体块；assistant/system 附件只回放文本清单，避免把图片放到供应商不接受的角色。 -SessionManager 负责组装会话上下文：系统提示、Skills、召回的 Knowledge、压缩后的 Timeline、可选的 active plan 摘要和当前消息历史。`session::turn_input` 在 Session 锁外并行读取 Knowledge、active plan 并压缩历史，然后通过同一个 assembly 路径生成首次请求和 context-overflow 重试输入；重试不得复制系统提示或 runtime context 拼接逻辑。普通闲聊 session 没有 plan 摘要；计划状态由 `WorkManager` 从 SQLite 读取，因此不以自然语言摘要作为权威来源。`AgentLoop` 接收完整输入执行一次模型/工具循环，本身不拥有会话状态。执行工具时额外传递只包含 session/turn 身份的 `ToolExecutionContext`；无状态工具使用默认实现忽略它，有状态外部适配器必须用它隔离资源，不能自行反向查询 SessionManager。 +SessionManager 负责组装会话上下文：系统提示、Skills、召回的 Knowledge、压缩后的 Timeline、可选的 active plan 摘要和当前消息历史。`session::turn_input` 在 Session 锁外并行读取 Knowledge、active plan 并压缩历史，然后通过同一个 assembly 路径生成首次请求和 context-overflow 重试输入；重试不得复制系统提示或 runtime context 拼接逻辑，也不得丢失已经从 mailbox 取出的 steering。普通闲聊 session 没有 plan 摘要；计划状态由 `WorkManager` 从 SQLite 读取，因此不以自然语言摘要作为权威来源。`AgentLoop` 接收完整输入执行一次模型/工具循环，本身不拥有会话状态，但通过本 Turn 的 mailbox 在安全边界接收追加用户输入。执行工具时额外传递只包含 session/turn 身份的 `ToolExecutionContext`；无状态工具使用默认实现忽略它，有状态外部适配器必须用它隔离资源，不能自行反向查询 SessionManager。 当前 Turn 的工具进度只从 `AgentLoop` 的结构化 `TurnEvent` 进入 `TurnController`，不能另建字符串 notification 通道重复投递。后台子 Agent 的 `TaskNotification` 表达跨 Turn 的任务完成，仍由独立的受监督消费者投递。自动标题属于非关键派生工作：Turn 持久化完成后由 `TaskSupervisor` 调度，Session worker 不等待模型生成；同一 Session 同时最多有一个标题任务，提交时仍校验标题保持默认值，避免覆盖用户改名。 @@ -251,7 +260,7 @@ Gateway 在 `/` 提供随二进制编译的 HTML/CSS/JavaScript，不依赖外 WebUI/TUI 文件字节通过受鉴权的 HTTP 接口流式传输，WebSocket 只携带短期 `upload_id` 和结构化附件描述。`UploadRegistry` 在内存中按 `cli_chat` chat scope 校验并消费待发送上传；消息继续以 `media_refs` 保存 Gateway 本地路径，不建立永久附件资产。下载接口必须通过 client、session、message 和附件序号反查路径，不能接受客户端路径。历史附件路径失效属于正常状态，不得影响历史文本读取。待发送但未进入消息的上传由 `TaskSupervisor` 所有的限时清理任务回收。Agent 上下文会为所有媒体注入内部路径清单，客户端响应不得暴露该路径。 -工具默认通过 `ToolResult` 返回文本；需要把图片等产物交给模型时，通过 `Tool::execute_with_media` 返回文本和结构化 `MediaRef`。工具只负责经过自身路径策略校验后声明媒体，不感知当前模型或 Provider。`AgentLoop` 仅将最新连续工具结果批次的媒体交给 `MediaHandlerRegistry`，旧工具媒体只回放文本和路径，避免历史 Base64 膨胀。OpenAI-compatible Provider 保持 `tool` 结果为文本，并在完整工具批次后构造仅存在于请求内的临时多模态 `user` 消息；Anthropic Provider 将媒体放入对应 `tool_result.content`。媒体加载、格式或能力检查失败必须降级成文本，不得使历史记录不可读取。 +工具默认通过 `ToolResult` 返回文本；需要把图片等产物交给模型时，通过 `Tool::execute_with_media` 返回文本和结构化 `MediaRef`。工具只负责经过自身路径策略校验后声明媒体，不感知当前模型或 Provider。`AgentLoop` 仅回放最新工具调用批次的原生媒体；紧随工具结果的 steering 不会使该批次的图片失去可见性，而后续 assistant 回复会终止其原生媒体回放，避免历史 Base64 膨胀。OpenAI-compatible Provider 保持 `assistant(tool_calls) → tool results → user steering` 的原生顺序；Anthropic Provider 把同批 `tool_result` 与紧随其后的 user steering 合并为一个 API 所需的 `role=user` 内容数组，持久化消息仍彼此独立。媒体加载、格式或能力检查失败必须降级成文本，不得使历史记录不可读取。 `browser` 是有状态工具适配器：`BrowserTool` 保持模型侧 action schema，`BrowserManager` 把 PicoBot dialog 映射到随机 agent-browser session，并用每 session mutex 保证同一页面串行、不同 dialog 并发；`AgentBrowserRunner` 以 argv 和 `--json` 调用外部原生 CLI，设置硬超时、输出/content boundaries/domain allowlist，底层 daemon 通过 Chrome CDP 工作。PicoBot 不链接 agent-browser 内部 crate、不直接暴露其 MCP、不使用 Fantoccini/ChromeDriver/WebDriver。截图只能写入配置的 artifact directory，并经 `ToolResultWithMedia` 返回。完整边界见 [AGENT_BROWSER_INTEGRATION.md](AGENT_BROWSER_INTEGRATION.md)。 diff --git a/resources/skills/about-picobot/references/architecture.md b/resources/skills/about-picobot/references/architecture.md index 96c2c25..7a79c53 100644 --- a/resources/skills/about-picobot/references/architecture.md +++ b/resources/skills/about-picobot/references/architecture.md @@ -42,9 +42,9 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message - Channels 通过 MessageBus 发布入站消息，通过 OutboundDispatcher 或每 Turn 一个的 TurnSink 接收出站写入，不感知 session 或 LLM - MessageBus 本体持有三条有界队列；出站路由、顺序和重试由 `OutboundDispatcher` 负责 -- SessionManager 拥有 session 状态、dialog 路由、上下文构建和每 session worker，并通过 worker 创建 AgentLoop +- SessionManager 拥有 session 状态、dialog 路由、上下文构建、每 session worker 和活动 Turn 的 steering mailbox，并通过 worker 创建 AgentLoop - TurnController 是活动 Turn 状态的唯一 owner；Session 在消息原子提交成功后才发布 Completed -- AgentLoop 跨轮无状态，接收已准备的 history 调用 LLM、执行工具并返回一次结果 +- AgentLoop 跨轮无状态，接收已准备的 history，并在安全模型边界排空本 Turn steering 后调用 LLM、执行工具并返回一次结果 - Providers 是纯 HTTP 流客户端，无 bus/session/channel 感知；签名 reasoning 状态只回放给匹配 Provider，不下发客户端或 Channel - DeliveryCoordinator 只投影完整快照，不修改会话历史；慢消费者跳过中间 revision，终态显式、有界投递 - 每个活动 Turn 独占一个 TurnSink；平台 message ID 和 reaction 清理状态只存在于 sink 内 @@ -64,7 +64,8 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message - OutboundDispatcher 通过 ChannelManager 路由出站消息 - 配置目录 `.env` 与 workspace `.env` 仅在单线程启动阶段分层加载，并使用 `unsafe { env::set_var(...) }` 写入进程环境；优先级为既有进程环境 > workspace > 配置目录 - `browser` 工具默认启用，只有 `browser.enabled=false` 时不注册；缺少 CLI/Chrome 不阻止 Gateway 启动，但 health 和实际调用会给出安装错误。每个 PicoBot dialog 映射到独立 agent-browser session，底层原生 daemon 使用 Chrome CDP，不依赖 Fantoccini/ChromeDriver/WebDriver -- 同一 session 的普通消息串行处理，不同 session 可并发；session 队列容量为 32，满时明确拒绝 +- 同一 session 只运行一个 Turn；活动 Turn 期间普通输入默认 steering，`/queue` 明确等待下一 Turn，不同 session 可并发 +- steering mailbox 容量为 32 条/64 KiB，满或关闭时可靠回退到容量 32 的 session 队列；两者都无法接收时明确拒绝 - 出站消息按 `(channel, chat_id)` 分 lane 保序；lane 容量为 64，慢目标不阻塞其他目标 - 活动 Turn 与普通出站消息共享 `(channel, chat_id)` 写锁；禁止把 token delta 放入 MessageBus - `cli_chat` 向 TUI/WebUI 发送统一 `turn_updated` 完整快照；飞书默认 FinalOnly，开启 `live_updates` 后编辑同一卡片 @@ -138,7 +139,7 @@ create → 存入 Storage → 载入 memory → 设为当前 dialog ### 消息处理与并发 -普通消息先 `try_send` 到该 session 的有界 worker 队列，Gateway 主 processor 随即返回 `AgentProcessing`。Slash command 直接执行，不进入此队列，因此 `/stop` 不会排在长模型调用后。 +没有活动 Turn 时，普通消息先 `try_send` 到该 session 的有界 worker 队列，Gateway 主 processor 随即返回 `AgentProcessing`。活动 Turn 期间普通消息默认进入有界 steering mailbox，并在完整工具批次结束后或无工具最终回复边界作为真实 `role=user` 消息注入下一次模型调用；`/queue ` 绕过 mailbox，明确进入下一 Turn。`/stop` 直接取消当前 Turn 并清空 mailbox 与普通队列，不会排在长模型调用后。 Worker 的处理原则： @@ -147,6 +148,7 @@ Worker 的处理原则： 3. 提交由旧快照产生的结果前重新验证 generation/version，防止 `/stop`、`/clear` 或 `/delete` 后写回陈旧状态。 4. Session 持久化由独立 `persistence_lock` 串行化；批量消息使用原子写入，失败时精确回滚内存后缀。 5. 上下文溢出时按 Provider 返回的真实限制重新压缩并重试。 +6. mailbox 的接收与关闭原子互斥；所有输入在 Session 锁内取得单调序号，未消费 steering 与普通队列按该序号恢复，不能丢失或互相超越。 WebUI/TUI 的 Active Turn 使用 `send_message(files=...)` 向自身 session 投递附件时，附件暂存到 task-local Turn delivery，成功结束后并入最终 assistant 消息，因此工具链始终排在附件回复之前且不会出现自引用来源前缀。其他自投递要求 task-local Turn ID 与 session 的 active Turn 匹配；历史中的 assistant/system 附件只作为文本清单提供给模型，原生媒体块仅用于 user 输入和当前工具结果。 @@ -299,4 +301,5 @@ Gateway 关停顺序： | `/?`, `/help` | 显示帮助 | | `/mcp` | 显示 MCP 状态 | | `/health` | 检查 PicoBot 运行依赖 | +| `/queue ` | 等当前 Turn 完成后作为下一 Turn 处理 | | `/stop` | 停止当前任务并清空消息队列 | diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 57fd9aa..5e2b5a0 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1,5 +1,6 @@ use crate::agent::context_compressor::estimate_tokens; use crate::agent::media_handler::MediaHandlerRegistry; +use crate::agent::steering::SteeringDrain; use crate::agent::system_prompt::build_system_prompt; use crate::agent::turn_event::{AgentTurnContext, TurnEvent}; use crate::bus::message::ContentBlock; @@ -37,11 +38,33 @@ fn should_include_message_media(messages: &[ChatMessage], index: usize) -> bool if message.role != "tool" { return true; } - let active_tool_start = messages + + // Tool media is replayed only for the most recent tool-call batch. A + // steering message is a real `user` message and therefore breaks the + // contiguous `assistant(tool_calls), tool...` shape used by older code. + // Find the latest assistant tool-call declaration instead, then keep its + // tool results eligible until the next assistant message. This keeps a + // screenshot returned by a tool visible when the user steers immediately + // after that tool batch. + let Some(tool_call_start) = messages.iter().rposition(|candidate| { + candidate.role == "assistant" + && candidate + .tool_calls + .as_ref() + .is_some_and(|calls| !calls.is_empty()) + }) else { + return false; + }; + if index <= tool_call_start { + return false; + } + let no_assistant_before = !messages[tool_call_start + 1..index] .iter() - .rposition(|candidate| candidate.role != "tool") - .map_or(0, |last_non_tool| last_non_tool + 1); - index >= active_tool_start + .any(|candidate| candidate.role == "assistant"); + let no_assistant_after = !messages[index + 1..] + .iter() + .any(|candidate| candidate.role == "assistant"); + no_assistant_before && no_assistant_after } /// Build content blocks from text and media, respecting model input capabilities @@ -572,6 +595,50 @@ impl AgentLoop { } } + /// Add steering messages to the in-memory transcript in receive order. + /// They remain ordinary `role=user` messages so every provider sees the + /// same conversation semantics and persistence can commit them alongside + /// the rest of this turn. + fn append_steering_messages( + messages: &mut Vec, + emitted_messages: &mut Vec, + consumed_steering: &mut Vec, + steering_messages: Vec, + turn: &AgentTurnContext, + iteration: u32, + ) { + for mut message in steering_messages { + // Session routes only ordinary user input to the mailbox. Keep a + // defensive normalization here because the mailbox is public and + // can also be used by embedders/tests. + message.role = "user".to_string(); + if message.turn_id.is_none() { + message.turn_id = Some(turn.turn_id.clone()); + } + if message.iteration.is_none() { + message.iteration = Some(iteration); + } + consumed_steering.push(message.clone()); + messages.push(message.clone()); + emitted_messages.push(message); + } + } + + fn close_steering(turn: Option<&AgentTurnContext>) { + if let Some(mailbox) = turn.and_then(AgentTurnContext::steering) { + mailbox.close(); + } + } + + fn restore_steering(turn: Option<&AgentTurnContext>, consumed_steering: Vec) { + if consumed_steering.is_empty() { + return; + } + if let Some(mailbox) = turn.and_then(AgentTurnContext::steering) { + mailbox.restore_front(consumed_steering); + } + } + fn chat_message_to_llm_message(&self, m: &ChatMessage, include_media: bool) -> Message { let content = if m.media_refs.is_empty() || !include_media { vec![ContentBlock::text(&m.content)] @@ -693,6 +760,11 @@ impl AgentLoop { // Track tool calls for loop detection let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default()); let mut emitted_messages = Vec::new(); + // Steering messages are removed from the mailbox only at safe + // boundaries. Keep a local copy until this invocation commits; if a + // later provider/tool request fails, restore them before Session + // retries from persisted history. + let mut consumed_steering = Vec::new(); let mut accumulated_tokens: u32 = 0; let mut accumulated_usage = crate::providers::Usage::default(); let mut last_request_usage = None; @@ -700,6 +772,7 @@ impl AgentLoop { for iteration in 0..self.max_iterations { #[cfg(debug_assertions)] tracing::debug!(iteration, "Agent iteration started"); + let last_iteration = iteration.saturating_add(1) >= self.max_iterations; // Preemptive context check: trim old tool results if token estimate // exceeds 80% of context window to prevent mid-loop overflow. @@ -738,11 +811,27 @@ impl AgentLoop { }; // Call LLM - let iteration = u32::try_from(iteration) - .map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?; - let response = self + let iteration = match u32::try_from(iteration) { + Ok(iteration) => iteration, + Err(_) => { + Self::restore_steering(turn.as_ref(), consumed_steering); + return Err(AgentError::Other("tool iteration exceeds u32".to_string())); + } + }; + let response = match self .stream_completion(request, iteration, turn.as_ref()) - .await?; + .await + { + Ok(response) => response, + Err(error) => { + // The invocation may be retried from persisted history. + // Restore every steering message consumed by an earlier + // boundary; Session decides whether to retry or close + // and queue them after receiving this error. + Self::restore_steering(turn.as_ref(), consumed_steering); + return Err(error); + } + }; accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens); merge_usage(&mut accumulated_usage, &response.usage); @@ -756,11 +845,45 @@ impl AgentLoop { "LLM response received" ); - // If no tool calls, this is the final response + // If no tool calls, this is normally the final response. When a + // steering message arrived and another model iteration remains, + // preserve this assistant message as an intermediate transcript + // entry and continue with the user input. At the last iteration + // there is no budget for another normal request, so close the + // mailbox and let Session move any pending messages to its queue. if response.tool_calls.is_empty() { let mut assistant_message = ChatMessage::assistant(response.content); assistant_message.reasoning_content = response.reasoning_content; assistant_message.provider_state = response.provider_state; + + let steering = turn.as_ref().and_then(AgentTurnContext::steering); + let pending = if last_iteration { + if let Some(mailbox) = steering.as_ref() { + mailbox.close(); + } + None + } else { + steering.as_ref().map(|mailbox| mailbox.drain_or_close()) + }; + + if let Some(SteeringDrain::Messages(steering_messages)) = pending { + let Some(turn_context) = turn.as_ref() else { + unreachable!("steering messages require a turn context"); + }; + Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, false); + messages.push(assistant_message.clone()); + emitted_messages.push(assistant_message); + Self::append_steering_messages( + &mut messages, + &mut emitted_messages, + &mut consumed_steering, + steering_messages, + turn_context, + iteration, + ); + continue; + } + Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, true); emitted_messages.push(assistant_message.clone()); crate::observability::metrics::global_metrics().record_turn( @@ -776,10 +899,13 @@ impl AgentLoop { }); } - if let Some(turn) = turn.as_ref() { - turn.emitter + if let Some(turn) = turn.as_ref() + && let Err(error) = turn + .emitter .emit(TurnEvent::TextSegmentFinished { iteration }) - .map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?; + { + Self::restore_steering(Some(turn), consumed_steering); + return Err(AgentError::Other(format!("turn event rejected: {error}"))); } // Execute tool calls. User-visible progress is emitted through the @@ -809,14 +935,21 @@ impl AgentLoop { emitted_messages.push(assistant_message); // Execute tools and add results to messages - let tool_results = self + let tool_results = match self .execute_tools( &response.tool_calls, iteration, turn.as_ref(), &tool_context, ) - .await?; + .await + { + Ok(results) => results, + Err(error) => { + Self::restore_steering(turn.as_ref(), consumed_steering); + return Err(error); + } + }; for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) { // Log function call with name and arguments @@ -866,6 +999,31 @@ impl AgentLoop { } } + // A complete tool batch is the first safe steering boundary. Do + // not drain at the final available iteration: those inputs must + // remain in the closed mailbox for Session to queue after this + // turn rather than being silently consumed by the summary call. + if let Some(mailbox) = turn.as_ref().and_then(AgentTurnContext::steering) { + if last_iteration { + mailbox.close(); + } else { + let steering_messages = mailbox.drain(); + if !steering_messages.is_empty() { + let Some(turn_context) = turn.as_ref() else { + unreachable!("steering messages require a turn context"); + }; + Self::append_steering_messages( + &mut messages, + &mut emitted_messages, + &mut consumed_steering, + steering_messages, + turn_context, + iteration, + ); + } + } + } + // Loop continues to next iteration with updated messages #[cfg(debug_assertions)] tracing::debug!( @@ -876,6 +1034,10 @@ impl AgentLoop { } // Max iterations reached - ask LLM for a summary based on completed work + // Any mailbox input still pending at this boundary cannot be consumed + // without exceeding the configured tool-iteration budget. Keep it in + // the closed mailbox for Session's next-turn fallback. + Self::close_steering(turn.as_ref()); tracing::warn!("Max iterations reached, requesting final summary from LLM"); // Add a message asking for summary @@ -895,8 +1057,13 @@ impl AgentLoop { tools: None, // No tools in final summary call }; - let summary_iteration = u32::try_from(self.max_iterations) - .map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?; + let summary_iteration = match u32::try_from(self.max_iterations) { + Ok(iteration) => iteration, + Err(_) => { + Self::restore_steering(turn.as_ref(), consumed_steering); + return Err(AgentError::Other("tool iteration exceeds u32".to_string())); + } + }; match self .stream_completion(request, summary_iteration, turn.as_ref()) .await @@ -934,8 +1101,9 @@ impl AgentLoop { "I reached the maximum number of tool call iterations ({}) without completing the task. The work done so far has been lost due to an error. Please try breaking the task into smaller steps.", self.max_iterations ); - if let Some(turn) = turn.as_ref() { - turn.emitter + if let Some(turn) = turn.as_ref() + && let Err(error) = turn + .emitter .emit(TurnEvent::TextSegmentFinished { iteration: summary_iteration, }) @@ -945,9 +1113,9 @@ impl AgentLoop { delta: fallback.clone(), }) }) - .map_err(|error| { - AgentError::Other(format!("turn event rejected: {error}")) - })?; + { + Self::restore_steering(Some(turn), consumed_steering); + return Err(AgentError::Other(format!("turn event rejected: {error}"))); } let mut final_message = ChatMessage::assistant(fallback); Self::annotate_message(&mut final_message, turn.as_ref(), summary_iteration, true); @@ -1154,6 +1322,7 @@ impl AgentLoop { #[cfg(test)] mod tests { use super::*; + use crate::agent::SteeringMailbox; use crate::observability::{MultiObserver, Observer}; use crate::providers::{ ChatCompletionResponse, FinishReason, ProviderChunk, ProviderStream, Usage, @@ -1167,6 +1336,10 @@ mod tests { struct StreamingTextProvider; + struct ErrorAfterFirstProvider { + requests: std::sync::Mutex, + } + #[async_trait::async_trait] impl LLMProvider for StreamingTextProvider { async fn stream( @@ -1212,6 +1385,54 @@ mod tests { } } + #[async_trait::async_trait] + impl LLMProvider for ErrorAfterFirstProvider { + async fn stream( + &self, + _request: ChatCompletionRequest, + ) -> Result { + let request_number = { + let mut requests = self.requests.lock().unwrap(); + *requests += 1; + *requests + }; + if request_number > 1 { + return Err(Box::new(std::io::Error::other( + "synthetic provider failure", + ))); + } + let chunks = vec![ + ProviderChunk::Metadata { + id: "first".into(), + model: "error-after-first".into(), + }, + ProviderChunk::Text("first response".into()), + ProviderChunk::Usage(Usage { + prompt_tokens: 1, + completion_tokens: 1, + total_tokens: 2, + ..Usage::default() + }), + ProviderChunk::Done(FinishReason::Stop), + ]; + Ok(Box::pin(futures_util::stream::iter( + chunks.into_iter().map(Ok), + ))) + } + + fn ptype(&self) -> &str { + "test" + } + + fn name(&self) -> &str { + "error-after-first" + } + + fn model_id(&self) -> &str { + "error-after-first" + } + } + #[tokio::test] async fn process_streaming_emits_turn_blocks_and_stamps_durable_message() { let agent = AgentLoop::with_provider( @@ -1264,6 +1485,87 @@ mod tests { )); } + #[tokio::test] + async fn steering_after_final_response_is_injected_at_next_model_boundary() { + let agent = AgentLoop::with_provider( + Arc::new(StreamingTextProvider), + 2, + "streaming-test".into(), + PathBuf::from("."), + Vec::new(), + ); + let mailbox = Arc::new(SteeringMailbox::new()); + mailbox + .try_push(ChatMessage::user("please include the log summary")) + .unwrap(); + let (controller, emitter, _) = TurnController::start("session", "assistant-id"); + let initial = controller.snapshot(); + let context = AgentTurnContext::new_with_steering( + initial.id.0.clone(), + initial.message_id.clone(), + emitter, + mailbox.clone(), + ); + + let result = agent + .process_streaming(vec![ChatMessage::user("hi")], context) + .await + .unwrap(); + + let positions: Vec<(usize, &str)> = result + .emitted_messages + .iter() + .enumerate() + .map(|(index, message)| (index, message.role.as_str())) + .collect(); + assert_eq!( + positions, + vec![(0, "assistant"), (1, "user"), (2, "assistant")] + ); + assert_eq!( + result.emitted_messages[1].content, + "please include the log summary" + ); + assert_eq!(result.final_response.content, "hello world"); + assert!(mailbox.is_closed()); + assert!(mailbox.is_empty()); + } + + #[tokio::test] + async fn provider_error_after_steering_drain_restores_input_for_retry() { + let agent = AgentLoop::with_provider( + Arc::new(ErrorAfterFirstProvider { + requests: std::sync::Mutex::new(0), + }), + 2, + "error-after-first".into(), + PathBuf::from("."), + Vec::new(), + ); + let mailbox = Arc::new(SteeringMailbox::new()); + mailbox.try_push(ChatMessage::user("retry me")).unwrap(); + let (controller, emitter, _) = TurnController::start("session", "assistant-id"); + let initial = controller.snapshot(); + let context = AgentTurnContext::new_with_steering( + initial.id.0.clone(), + initial.message_id.clone(), + emitter, + mailbox.clone(), + ); + + let error = agent + .process_streaming(vec![ChatMessage::user("hi")], context) + .await + .unwrap_err(); + assert!( + matches!(error, AgentError::LlmError(message) if message.contains("synthetic provider failure")) + ); + assert!(!mailbox.is_closed()); + let restored = mailbox.take_pending(); + assert_eq!(restored.len(), 1); + assert_eq!(restored[0].content, "retry me"); + } + impl TestObserver { fn new() -> Self { Self { @@ -1420,6 +1722,69 @@ mod tests { ))); } + #[tokio::test] + async fn tool_media_remains_visible_when_steering_follows_tool_batch() { + use std::io::Write; + + let mut image = tempfile::Builder::new().suffix(".png").tempfile().unwrap(); + image.write_all(b"\x89PNG\r\n\x1a\nminimal").unwrap(); + let provider = Arc::new(ToolMediaProvider { + image_path: image.path().to_string_lossy().into_owned(), + requests: std::sync::Mutex::new(Vec::new()), + }); + let tools = Arc::new(ToolRegistry::new()); + tools.register(FileReadTool::new()); + let agent = AgentLoop::with_provider_and_tools( + provider.clone(), + tools, + 2, + "vision-test".to_string(), + std::env::current_dir().unwrap(), + vec!["text".to_string(), "image".to_string()], + ); + + let mailbox = Arc::new(SteeringMailbox::new()); + mailbox + .try_push(ChatMessage::user("also explain what you found")) + .unwrap(); + let (controller, emitter, _) = TurnController::start("session", "assistant-id"); + let turn = controller.snapshot(); + let context = AgentTurnContext::new_with_steering( + turn.id.0.clone(), + turn.message_id.clone(), + emitter, + mailbox, + ); + let result = agent + .process_streaming(vec![ChatMessage::user("inspect the image")], context) + .await + .unwrap(); + + assert_eq!(result.final_response.content, "image seen"); + let requests = provider.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + let second_messages = &requests[1].messages; + let tool_index = second_messages + .iter() + .position(|message| message.role == "tool") + .unwrap(); + assert!( + second_messages[tool_index] + .content + .iter() + .any(|block| matches!(block, ContentBlock::ImageUrl { .. })) + ); + assert!( + second_messages[tool_index + 1..] + .iter() + .any(|message| message.role == "user" + && message.content.iter().any(|block| matches!( + block, + ContentBlock::Text { text } if text.contains("also explain what you found") + ))) + ); + } + #[test] fn test_should_execute_in_parallel_single_tool() { // Would need a proper setup with AgentLoop to test fully @@ -1658,6 +2023,11 @@ mod tests { assert!(should_include_message_media(&messages, 4)); messages.push(ChatMessage::user("next turn")); + // A same-turn steering message is a real user message, but it must + // not hide media returned by the immediately preceding tool batch. + assert!(should_include_message_media(&messages, 4)); + + messages.push(ChatMessage::assistant("final")); assert!(!should_include_message_media(&messages, 4)); } } diff --git a/src/agent/mod.rs b/src/agent/mod.rs index dc75e49..1ed26e5 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -1,12 +1,14 @@ pub mod agent_loop; pub mod context_compressor; pub mod media_handler; +pub mod steering; pub mod sub_agent; pub mod system_prompt; pub mod turn_event; pub use agent_loop::{AgentError, AgentLoop, AgentProcessResult}; pub use context_compressor::{ContextCompressor, estimate_tokens}; +pub use steering::{SteeringDrain, SteeringMailbox, SteeringPushError}; pub use sub_agent::{ DelegateContext, ExecutionMode, SubAgentConfig, SubAgentError, SubAgentManager, SubAgentResult, TaskNotification, TaskStatus, diff --git a/src/agent/steering.rs b/src/agent/steering.rs new file mode 100644 index 0000000..e9b30b1 --- /dev/null +++ b/src/agent/steering.rs @@ -0,0 +1,489 @@ +//! Bounded, session-owned mailbox for same-turn user steering. +//! +//! A mailbox is intentionally separate from the session work queue. The +//! gateway can accept a normal user message while a turn is running and place +//! it here; [`AgentLoop`](super::AgentLoop) drains it only at safe model +//! boundaries (after a complete tool batch, or before deciding that a +//! response is final). The state transition performed by +//! [`SteeringMailbox::drain_or_close`] is atomic with respect to producers, +//! which means an input is either accepted by the active turn or rejected so +//! the caller can put it on the next-turn queue -- never both and never +//! neither. + +use crate::bus::ChatMessage; +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +/// Default maximum number of steering messages accepted by one active turn. +pub const DEFAULT_MAX_STEERING_MESSAGES: usize = 32; +/// Default aggregate UTF-8 byte budget for pending steering messages. +pub const DEFAULT_MAX_STEERING_BYTES: usize = 64 * 1024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MailboxPhase { + Accepting, + Closed, +} + +#[derive(Debug)] +struct MailboxState { + phase: MailboxPhase, + pending: VecDeque, + pending_bytes: usize, + /// Messages drained at a safe boundary but not yet committed to durable + /// history. Keeping their count/size reserved prevents concurrent + /// producers from filling the capacity that an error retry may need to + /// restore. + in_flight_messages: usize, + in_flight_bytes: usize, + /// Exact drained messages retained until commit. This lets Session + /// recover a successful AgentLoop result if its subsequent persistence + /// transaction fails. + in_flight: VecDeque, +} + +/// Error returned when the active turn cannot accept a steering message. +/// +/// The original message is returned in the error so the caller can enqueue it +/// as a normal next-turn task without cloning or losing media metadata. +#[derive(Debug, Clone)] +pub enum SteeringPushError { + /// The turn has reached a terminal boundary. Route the message to the + /// session's ordinary queue. + Closed(Box), + /// The mailbox is accepting input, but its bounded capacity is exhausted. + /// Route the message to the ordinary queue (and normally notify the user). + Full(Box), +} + +impl SteeringPushError { + /// Recover the message that was rejected by [`SteeringMailbox::try_push`]. + pub fn into_message(self) -> ChatMessage { + match self { + Self::Closed(message) | Self::Full(message) => *message, + } + } + + pub fn is_closed(&self) -> bool { + matches!(self, Self::Closed(_)) + } + + pub fn is_full(&self) -> bool { + matches!(self, Self::Full(_)) + } +} + +/// Result of the atomic final-response boundary operation. +#[derive(Debug, Clone)] +pub enum SteeringDrain { + /// One or more inputs were accepted and removed from the mailbox. The + /// mailbox remains open for a subsequent safe boundary. + Messages(Vec), + /// No pending input existed. The mailbox is now closed; later producers + /// receive [`SteeringPushError::Closed`]. + Closed, +} + +/// Shared state for user steering during one active AgentLoop execution. +/// +/// Cloning a mailbox is cheap and shares the same mutex-protected state. In +/// practice the session stores an `Arc` in its active-turn +/// handle and gives another clone to [`AgentTurnContext`](super::AgentTurnContext). +#[derive(Clone)] +pub struct SteeringMailbox { + state: Arc>, + max_messages: usize, + max_bytes: usize, +} + +impl std::fmt::Debug for SteeringMailbox { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + formatter + .debug_struct("SteeringMailbox") + .field("phase", &state.phase) + .field("pending_messages", &state.pending.len()) + .field("pending_bytes", &state.pending_bytes) + .field("max_messages", &self.max_messages) + .field("max_bytes", &self.max_bytes) + .finish() + } +} + +impl SteeringMailbox { + /// Construct a mailbox using the product defaults (32 messages/64 KiB). + pub fn new() -> Self { + Self::with_limits(DEFAULT_MAX_STEERING_MESSAGES, DEFAULT_MAX_STEERING_BYTES) + } + + /// Construct a mailbox with explicit bounded capacities. Zero limits are + /// allowed and make every push return [`SteeringPushError::Full`]. + pub fn with_limits(max_messages: usize, max_bytes: usize) -> Self { + Self { + state: Arc::new(Mutex::new(MailboxState { + phase: MailboxPhase::Accepting, + pending: VecDeque::new(), + pending_bytes: 0, + in_flight_messages: 0, + in_flight_bytes: 0, + in_flight: VecDeque::new(), + })), + max_messages, + max_bytes, + } + } + + /// Return an `Arc` suitable for storing in Session and AgentTurnContext. + pub fn new_shared() -> Arc { + Arc::new(Self::new()) + } + + /// Return an `Arc` suitable for storing in Session with explicit limits. + pub fn shared_with_limits(max_messages: usize, max_bytes: usize) -> Arc { + Arc::new(Self::with_limits(max_messages, max_bytes)) + } + + /// Try to accept one real user [`ChatMessage`]. + /// + /// This operation and the final close operation use the same mutex. A + /// producer racing with `drain_or_close` therefore receives a deterministic + /// result and can route a rejected message to the ordinary queue. + pub fn try_push(&self, message: ChatMessage) -> Result<(), SteeringPushError> { + let message_bytes = message_size_bytes(&message); + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if state.phase == MailboxPhase::Closed { + return Err(SteeringPushError::Closed(Box::new(message))); + } + if state.pending.len().saturating_add(state.in_flight_messages) >= self.max_messages + || state + .pending_bytes + .saturating_add(state.in_flight_bytes) + .saturating_add(message_bytes) + > self.max_bytes + { + return Err(SteeringPushError::Full(Box::new(message))); + } + state.pending_bytes = state.pending_bytes.saturating_add(message_bytes); + state.pending.push_back(message); + Ok(()) + } + + /// Drain currently pending inputs while leaving the mailbox open. + /// + /// This is used after a complete tool-call batch. It intentionally does + /// not close the mailbox: another input may steer a later iteration. + pub fn drain(&self) -> Vec { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + drain_pending_locked(&mut state) + } + + /// Atomically drain pending inputs, or close the mailbox if it is empty. + pub fn drain_or_close(&self) -> SteeringDrain { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if state.pending.is_empty() { + state.phase = MailboxPhase::Closed; + SteeringDrain::Closed + } else { + SteeringDrain::Messages(drain_pending_locked(&mut state)) + } + } + + /// Close acceptance without dropping pending messages. Session uses + /// [`take_pending`](Self::take_pending) after AgentLoop returns to move + /// those messages to the ordinary next-turn queue (for example when the + /// maximum iteration budget is exhausted). + pub fn close(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.phase = MailboxPhase::Closed; + } + + /// Close acceptance and return all pending messages. This is convenient + /// for cancellation/error paths where the caller immediately owns the + /// rejected messages. Any in-flight batch is intentionally discarded; + /// `/stop` uses this method to preserve its queue-clearing semantics. + pub fn close_and_take_pending(&self) -> Vec { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.phase = MailboxPhase::Closed; + // Cancellation is an explicit discard boundary. Any in-flight + // messages that were already drained belong to this cancelled turn + // and must not reserve capacity forever. + take_pending_locked(&mut state) + } + + /// Restore all messages drained by AgentLoop since the last commit. This + /// is useful when the AgentLoop completed but Session's durable write then + /// failed: the next retry/queue operation can replay the exact accepted + /// inputs instead of silently losing them. + pub fn restore_drained(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + restore_in_flight_locked(&mut state); + } + + /// Restore messages drained by AgentLoop when a provider/tool error makes + /// the current invocation retry from persisted history. The messages are + /// prepended in their original order and their reserved capacity is + /// released. `drain()`/`drain_or_close()` reserve capacity while a batch is + /// in-flight, so this operation cannot overflow a bounded mailbox due to a + /// racing producer. + pub fn restore_front(&self, messages: Vec) { + if messages.is_empty() { + return; + } + let restored_bytes = messages.iter().map(message_size_bytes).sum::(); + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.in_flight_messages = state.in_flight_messages.saturating_sub(messages.len()); + state.in_flight_bytes = state.in_flight_bytes.saturating_sub(restored_bytes); + for _ in 0..messages.len() { + state.in_flight.pop_front(); + } + for message in messages.into_iter().rev() { + state.pending.push_front(message); + } + state.pending_bytes = state.pending_bytes.saturating_add(restored_bytes); + } + + /// Mark all previously drained messages as durably committed. Session + /// calls this only after the complete Turn persistence transaction + /// succeeds. It is a no-op when no steering batch was consumed. + pub fn commit_drained(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.in_flight_messages = 0; + state.in_flight_bytes = 0; + state.in_flight.clear(); + } + + /// Take pending inputs without changing whether producers may still push. + /// + /// Normally used after `close()`; keeping this method explicit makes it + /// possible for Session to transfer accepted-but-unprocessed input to its + /// FIFO queue without opening a race with a new turn. + pub fn take_pending(&self) -> Vec { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + take_pending_locked(&mut state) + } + + pub fn is_closed(&self) -> bool { + let state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.phase == MailboxPhase::Closed + } + + pub fn len(&self) -> usize { + let state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.pending.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn max_messages(&self) -> usize { + self.max_messages + } + + pub fn max_bytes(&self) -> usize { + self.max_bytes + } +} + +impl Default for SteeringMailbox { + fn default() -> Self { + Self::new() + } +} + +/// Approximate the bounded payload size without serializing the complete +/// message. Content, media paths/types, tool metadata and source fields are +/// all untrusted input; counting their UTF-8 bytes gives a conservative enough +/// guard while retaining the original message losslessly. +fn message_size_bytes(message: &ChatMessage) -> usize { + let mut bytes = message.id.len() + + message.role.len() + + message.content.len() + + message.reasoning_content.as_deref().map_or(0, str::len) + + message.turn_id.as_deref().map_or(0, str::len) + + message.tool_call_id.as_deref().map_or(0, str::len) + + message.tool_name.as_deref().map_or(0, str::len); + for media in &message.media_refs { + bytes = bytes.saturating_add(media.path.len() + media.media_type.len()); + } + if let Some(tool_calls) = &message.tool_calls { + for call in tool_calls { + bytes = bytes + .saturating_add(call.id.len()) + .saturating_add(call.name.len()) + .saturating_add(call.arguments.to_string().len()); + } + } + bytes +} + +fn drain_pending_locked(state: &mut MailboxState) -> Vec { + let messages: Vec<_> = state.pending.drain(..).collect(); + let bytes = messages.iter().map(message_size_bytes).sum::(); + state.pending_bytes = state.pending_bytes.saturating_sub(bytes); + state.in_flight_messages = state.in_flight_messages.saturating_add(messages.len()); + state.in_flight_bytes = state.in_flight_bytes.saturating_add(bytes); + state.in_flight.extend(messages.iter().cloned()); + messages +} + +fn take_pending_locked(state: &mut MailboxState) -> Vec { + state.pending_bytes = 0; + state.in_flight_messages = 0; + state.in_flight_bytes = 0; + state.in_flight.clear(); + state.pending.drain(..).collect() +} + +fn restore_in_flight_locked(state: &mut MailboxState) { + if state.in_flight.is_empty() { + return; + } + let messages: Vec<_> = state.in_flight.drain(..).collect(); + let bytes = messages.iter().map(message_size_bytes).sum::(); + state.in_flight_messages = 0; + state.in_flight_bytes = 0; + for message in messages.into_iter().rev() { + state.pending.push_front(message); + } + state.pending_bytes = state.pending_bytes.saturating_add(bytes); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::thread; + + #[test] + fn accepts_fifo_messages_and_clone_shares_state() { + let mailbox = SteeringMailbox::with_limits(2, 100); + let clone = mailbox.clone(); + mailbox.try_push(ChatMessage::user("one")).unwrap(); + clone.try_push(ChatMessage::user("two")).unwrap(); + assert_eq!(mailbox.len(), 2); + let messages = mailbox.drain(); + assert_eq!( + messages + .iter() + .map(|m| m.content.as_str()) + .collect::>(), + ["one", "two"] + ); + assert!(!mailbox.is_closed()); + } + + #[test] + fn rejects_full_message_without_losing_it() { + let mailbox = SteeringMailbox::with_limits(1, 10_000); + mailbox.try_push(ChatMessage::user("first")).unwrap(); + let second = ChatMessage::user("second"); + let error = mailbox.try_push(second.clone()).unwrap_err(); + assert!(error.is_full()); + assert_eq!(error.into_message().content, second.content); + assert_eq!(mailbox.len(), 1); + } + + #[test] + fn byte_limit_is_bounded() { + let mailbox = SteeringMailbox::with_limits(8, 3); + let message = ChatMessage::user("four"); + assert!(matches!( + mailbox.try_push(message), + Err(SteeringPushError::Full(_)) + )); + } + + #[test] + fn drain_or_close_is_atomic_and_preserves_close_race_semantics() { + let mailbox = Arc::new(SteeringMailbox::new()); + let producer = mailbox.clone(); + let close_result = thread::spawn(move || producer.drain_or_close()) + .join() + .unwrap(); + assert!(matches!(close_result, SteeringDrain::Closed)); + let message = ChatMessage::user("late"); + assert!(matches!( + mailbox.try_push(message), + Err(SteeringPushError::Closed(_)) + )); + } + + #[test] + fn drain_or_close_drains_but_keeps_accepting_when_non_empty() { + let mailbox = SteeringMailbox::new(); + mailbox.try_push(ChatMessage::user("first")).unwrap(); + let result = mailbox.drain_or_close(); + assert!(matches!(result, SteeringDrain::Messages(_))); + assert!(!mailbox.is_closed()); + mailbox.try_push(ChatMessage::user("second")).unwrap(); + assert_eq!(mailbox.drain()[0].content, "second"); + } + + #[test] + fn close_keeps_pending_for_next_turn_fallback() { + let mailbox = SteeringMailbox::new(); + mailbox.try_push(ChatMessage::user("defer")).unwrap(); + mailbox.close(); + assert!(mailbox.try_push(ChatMessage::user("late")).is_err()); + assert_eq!(mailbox.take_pending()[0].content, "defer"); + assert!(mailbox.is_empty()); + } + + #[test] + fn drained_capacity_is_reserved_until_commit_or_restore() { + let mailbox = SteeringMailbox::with_limits(1, 10_000); + mailbox.try_push(ChatMessage::user("first")).unwrap(); + let drained = mailbox.drain(); + assert_eq!(drained.len(), 1); + assert!(mailbox.try_push(ChatMessage::user("second")).is_err()); + mailbox.restore_front(drained); + assert_eq!(mailbox.take_pending()[0].content, "first"); + + mailbox.try_push(ChatMessage::user("committed")).unwrap(); + let _ = mailbox.drain(); + mailbox.commit_drained(); + mailbox.try_push(ChatMessage::user("after commit")).unwrap(); + + let drained = mailbox.drain(); + assert_eq!(drained[0].content, "after commit"); + mailbox.restore_drained(); + assert_eq!(mailbox.take_pending()[0].content, "after commit"); + } +} diff --git a/src/agent/turn_event.rs b/src/agent/turn_event.rs index da10f92..5ee8095 100644 --- a/src/agent/turn_event.rs +++ b/src/agent/turn_event.rs @@ -2,6 +2,7 @@ use std::sync::{Arc, Mutex}; use thiserror::Error; +use crate::agent::steering::SteeringMailbox; use crate::providers::ToolCall; /// Presentation facts emitted while AgentLoop processes one model turn. @@ -58,6 +59,9 @@ pub struct AgentTurnContext { pub turn_id: String, pub message_id: String, pub emitter: TurnEmitter, + /// Same-turn user input accepted while this turn is active. Session owns + /// the mailbox lifecycle; AgentLoop only drains it at safe boundaries. + pub steering: Option>, } impl AgentTurnContext { @@ -70,8 +74,36 @@ impl AgentTurnContext { turn_id: turn_id.into(), message_id: message_id.into(), emitter, + steering: None, } } + + /// Construct a streaming context with a shared steering mailbox. + pub fn new_with_steering( + turn_id: impl Into, + message_id: impl Into, + emitter: TurnEmitter, + steering: Arc, + ) -> Self { + Self { + turn_id: turn_id.into(), + message_id: message_id.into(), + emitter, + steering: Some(steering), + } + } + + /// Attach a mailbox to an existing context. This builder keeps the old + /// `AgentTurnContext::new` call sites source-compatible. + pub fn with_steering(mut self, steering: Arc) -> Self { + self.steering = Some(steering); + self + } + + /// Return a clone of the shared mailbox, if steering is enabled. + pub fn steering(&self) -> Option> { + self.steering.clone() + } } impl TurnEmitter { diff --git a/src/bus/message.rs b/src/bus/message.rs index 9a92569..94bcf8b 100644 --- a/src/bus/message.rs +++ b/src/bus/message.rs @@ -404,6 +404,10 @@ pub struct InboundMessage { pub channel: String, pub sender_id: String, pub chat_id: String, + /// Client-provided id for optimistic UI reconciliation. Channel-owned + /// inputs that do not expose a client id leave this unset; the session + /// layer may generate a durable id when it accepts the message. + pub client_message_id: Option, pub content: String, pub received_at: i64, pub media: Vec, diff --git a/src/channels/cli_chat.rs b/src/channels/cli_chat.rs index ec63aff..cb8bbf5 100644 --- a/src/channels/cli_chat.rs +++ b/src/channels/cli_chat.rs @@ -179,6 +179,7 @@ impl CliChatChannel { WsInbound::UserInput { content, upload_ids, + client_message_id, chat_id, .. } => { @@ -187,8 +188,15 @@ impl CliChatChannel { if content.trim().is_empty() && upload_ids.is_empty() { return Err(ChannelError::Other("Message is empty".to_string())); } + // `/queue` is deliberately allowed to carry attachments: it + // is a message-routing directive whose payload remains a + // normal user input. Other slash commands still reject + // attachments because their handlers do not consume media. + let slash_allows_attachments = crate::channels::parse_slash_command(&content) + .is_some_and(|(name, _)| name.eq_ignore_ascii_case("queue")); if !upload_ids.is_empty() && crate::channels::parse_slash_command(&content).is_some() + && !slash_allows_attachments { return Err(ChannelError::Other( "Attachments cannot be sent with slash commands".to_string(), @@ -200,6 +208,18 @@ impl CliChatChannel { "Chat does not belong to this client".to_string(), )); } + let client_message_id = + client_message_id.and_then(|raw| match uuid::Uuid::parse_str(&raw) { + Ok(id) => Some(id.to_string()), + Err(error) => { + tracing::warn!( + client_message_id = %raw, + error = %error, + "Ignoring invalid client message id" + ); + None + } + }); let uploads = self .uploads .take_many(&client.chat_id, &upload_ids) @@ -210,6 +230,7 @@ impl CliChatChannel { channel: self.name().to_string(), sender_id: "cli".to_string(), chat_id: target_chat_id, + client_message_id, content, received_at: crate::bus::message::current_timestamp(), media, @@ -1055,8 +1076,9 @@ mod tests { .handle_ws_inbound( client, WsInbound::UserInput { - content: "处理附件".into(), + content: "/queue 处理附件".into(), upload_ids: vec!["upload-1".into()], + client_message_id: Some("550e8400-e29b-41d4-a716-446655440000".into()), channel: None, chat_id: None, sender_id: None, @@ -1069,6 +1091,42 @@ mod tests { assert_eq!(inbound.media.len(), 1); assert_eq!(inbound.media[0].path, "/tmp/report.pdf"); assert_eq!(inbound.media[0].media_type, "file"); + assert_eq!( + inbound.client_message_id.as_deref(), + Some("550e8400-e29b-41d4-a716-446655440000") + ); + assert_eq!(inbound.content, "/queue 处理附件"); + } + + #[tokio::test] + async fn invalid_client_message_id_is_ignored_at_channel_boundary() { + let channel = CliChatChannel::new(); + let bus = MessageBus::new(4); + channel.start(bus.clone()).await.unwrap(); + let (sender, _receiver) = mpsc::channel(1); + let client = Arc::new(Client { + sender, + chat_id: "client".into(), + current_session_id: Mutex::new(None), + }); + + channel + .handle_ws_inbound( + client, + WsInbound::UserInput { + content: "hello".into(), + upload_ids: Vec::new(), + client_message_id: Some("not-a-uuid".into()), + channel: None, + chat_id: None, + sender_id: None, + }, + ) + .await + .unwrap(); + + let inbound = bus.consume_inbound().await.unwrap(); + assert_eq!(inbound.client_message_id, None); } #[tokio::test] diff --git a/src/channels/feishu.rs b/src/channels/feishu.rs index 53a5936..f1ed04a 100644 --- a/src/channels/feishu.rs +++ b/src/channels/feishu.rs @@ -788,7 +788,7 @@ impl FeishuChannel { let result: UploadResp = serde_json::from_str(&body_text).map_err(|e| { ChannelError::Other(format!( "Parse upload response error: {} | body: {}", - e, &body_text + e, body_text )) })?; @@ -876,7 +876,7 @@ impl FeishuChannel { let result: UploadResp = serde_json::from_str(&body_text).map_err(|e| { ChannelError::Other(format!( "Parse upload response error: {} | body: {}", - e, &body_text + e, body_text )) })?; @@ -1419,6 +1419,7 @@ impl FeishuChannel { channel: "feishu".to_string(), sender_id: parsed.open_id.clone(), chat_id: parsed.chat_id.clone(), + client_message_id: None, content: parsed.content, received_at: crate::bus::message::current_timestamp(), media: parsed.media, diff --git a/src/client/oneshot.rs b/src/client/oneshot.rs index 8fdf365..c876583 100644 --- a/src/client/oneshot.rs +++ b/src/client/oneshot.rs @@ -114,6 +114,7 @@ pub async fn run_once( let input = WsInbound::UserInput { content: prompt, upload_ids: Vec::new(), + client_message_id: None, channel: None, chat_id: None, sender_id: None, @@ -193,6 +194,7 @@ where let stop = WsInbound::UserInput { content: "/stop".to_string(), upload_ids: Vec::new(), + client_message_id: None, channel: None, chat_id: None, sender_id: None, diff --git a/src/client/tui/event.rs b/src/client/tui/event.rs index a80b639..28333c4 100644 --- a/src/client/tui/event.rs +++ b/src/client/tui/event.rs @@ -206,6 +206,7 @@ async fn handle_input_key(app: &mut App, key: KeyEvent) { WsInbound::UserInput { content: input, upload_ids, + client_message_id: None, channel: None, // Session routing is owned by the server. A full session // id is not a chat id and must never be sent here. diff --git a/src/config/mod.rs b/src/config/mod.rs index d52ac9d..6064b6d 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -563,9 +563,9 @@ pub struct LLMProviderConfig { impl LLMProviderConfig { pub fn cost_of(&self, prompt_tokens: u32, completion_tokens: u32) -> Option { match (self.price_input_per_million, self.price_output_per_million) { - (Some(pi), Some(po)) => Some( - prompt_tokens as f64 / 1e6 * pi + completion_tokens as f64 / 1e6 * po, - ), + (Some(pi), Some(po)) => { + Some(prompt_tokens as f64 / 1e6 * pi + completion_tokens as f64 / 1e6 * po) + } _ => None, } } diff --git a/src/gateway/http.rs b/src/gateway/http.rs index f669275..7319968 100644 --- a/src/gateway/http.rs +++ b/src/gateway/http.rs @@ -92,11 +92,17 @@ const EMBEDDED_FONTS: &[(&str, &[u8])] = &[ ), ( "jetbrains-mono-400.woff2", - include_bytes!(concat!(env!("OUT_DIR"), "/webui/fonts/jetbrains-mono-400.woff2")), + include_bytes!(concat!( + env!("OUT_DIR"), + "/webui/fonts/jetbrains-mono-400.woff2" + )), ), ( "jetbrains-mono-700.woff2", - include_bytes!(concat!(env!("OUT_DIR"), "/webui/fonts/jetbrains-mono-700.woff2")), + include_bytes!(concat!( + env!("OUT_DIR"), + "/webui/fonts/jetbrains-mono-700.woff2" + )), ), ]; @@ -824,7 +830,11 @@ pub async fn get_tools(State(state): State>) -> Result = entries .into_iter() .map(|(name, tool)| { - let source = if name.contains("__") { "mcp" } else { "builtin" }; + let source = if name.contains("__") { + "mcp" + } else { + "builtin" + }; json!({ "name": name, "description": tool.description(), diff --git a/src/gateway/router.rs b/src/gateway/router.rs index f69a9f8..bc7829b 100644 --- a/src/gateway/router.rs +++ b/src/gateway/router.rs @@ -420,6 +420,7 @@ mod tests { channel: "test".to_string(), sender_id: "user".to_string(), chat_id: "chat".to_string(), + client_message_id: None, content: "hello".to_string(), received_at: 123, media: vec![], diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 4c0208f..cc8365e 100644 --- a/src/gateway/ws.rs +++ b/src/gateway/ws.rs @@ -174,11 +174,7 @@ async fn handle_logs_socket(ws: WebSocket, query: WsLogsQuery) { let mut rx = tx.subscribe(); let (mut ws_sender, mut ws_receiver) = ws.split(); - let min_level = query - .level - .as_deref() - .map(parse_min_level) - .unwrap_or(0); + let min_level = query.level.as_deref().map(parse_min_level).unwrap_or(0); let search = query .search .filter(|s| !s.is_empty()) diff --git a/src/logging.rs b/src/logging.rs index daf495d..b0ed609 100644 --- a/src/logging.rs +++ b/src/logging.rs @@ -4,9 +4,9 @@ use tokio::sync::broadcast; use tracing::field::{Field, Visit}; use tracing::{Event, Subscriber}; use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use tracing_subscriber::Layer; use tracing_subscriber::layer::Context; use tracing_subscriber::registry::LookupSpan; -use tracing_subscriber::Layer; use tracing_subscriber::{ EnvFilter, fmt, fmt::time::LocalTime, layer::SubscriberExt, util::SubscriberInitExt, }; diff --git a/src/observability/metrics.rs b/src/observability/metrics.rs index e100a38..99657bd 100644 --- a/src/observability/metrics.rs +++ b/src/observability/metrics.rs @@ -76,11 +76,15 @@ impl Metrics { pub fn record_turn(&self, usage: Option<&Usage>, latency_ms: u64) { self.turns.fetch_add(1, Relaxed); if let Some(u) = usage { - self.tokens_in.fetch_add(u64::from(u.prompt_tokens), Relaxed); + self.tokens_in + .fetch_add(u64::from(u.prompt_tokens), Relaxed); self.tokens_out .fetch_add(u64::from(u.completion_tokens), Relaxed); } - let mut q = self.turn_latencies.lock().unwrap_or_else(|e| e.into_inner()); + let mut q = self + .turn_latencies + .lock() + .unwrap_or_else(|e| e.into_inner()); q.push_back(latency_ms); while q.len() > WINDOW { q.pop_front(); @@ -143,7 +147,10 @@ impl Metrics { } pub fn snapshot(&self) -> MetricsSnapshot { - let latencies = self.turn_latencies.lock().unwrap_or_else(|e| e.into_inner()); + let latencies = self + .turn_latencies + .lock() + .unwrap_or_else(|e| e.into_inner()); let p95 = percentile_95(&latencies); let providers = self.providers.lock().unwrap_or_else(|e| e.into_inner()); let mut cost = 0.0; diff --git a/src/protocol.rs b/src/protocol.rs index 1741b3d..d8f470a 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -138,6 +138,10 @@ pub enum WsInbound { content: String, #[serde(default, skip_serializing_if = "Vec::is_empty")] upload_ids: Vec, + /// Stable id generated by the client for optimistic-message + /// reconciliation. It is optional for older clients and channels. + #[serde(default, skip_serializing_if = "Option::is_none")] + client_message_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] channel: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -348,6 +352,36 @@ mod tests { assert_eq!(value["messages"][0]["id"], "message"); } + #[test] + fn user_input_preserves_optional_client_message_id() { + let inbound = parse_inbound( + r#"{"type":"user_input","content":"hello","client_message_id":"client-1"}"#, + ) + .unwrap(); + match inbound { + WsInbound::UserInput { + client_message_id, + upload_ids, + .. + } => { + assert_eq!(client_message_id.as_deref(), Some("client-1")); + assert!(upload_ids.is_empty()); + } + other => panic!("unexpected frame: {other:?}"), + } + + let serialized = serialize_inbound(&WsInbound::UserInput { + content: "hello".to_string(), + upload_ids: Vec::new(), + client_message_id: Some("client-1".to_string()), + channel: None, + chat_id: None, + sender_id: None, + }) + .unwrap(); + assert!(serialized.contains(r#""client_message_id":"client-1""#)); + } + #[test] fn history_defaults_new_reasoning_fields_for_old_frames() { let message: HistoryMessage = serde_json::from_value(serde_json::json!({ diff --git a/src/providers/anthropic.rs b/src/providers/anthropic.rs index 8659156..3e8339b 100644 --- a/src/providers/anthropic.rs +++ b/src/providers/anthropic.rs @@ -148,43 +148,82 @@ struct AnthropicMessage { } fn convert_messages(messages: &[Message]) -> Vec { - messages - .iter() - .map(|message| { - let role = if message.role == "tool" { - "user".to_string() - } else { - message.role.clone() - }; - let content = if let Some(ref tool_call_id) = message.tool_call_id { - vec![serde_json::json!({ + let mut converted = Vec::with_capacity(messages.len()); + let mut index = 0; + + while index < messages.len() { + let message = &messages[index]; + + // Anthropic requires all tool results for one assistant tool-use turn + // to be carried in a single `role: user` content array. Steering is + // represented as a normal user message in PicoBot history, so merge + // any immediately-following user messages into that same array at + // the provider boundary. Durable messages remain independent. + if message.role == "tool" && message.tool_call_id.is_some() { + let mut content = Vec::new(); + while index < messages.len() + && messages[index].role == "tool" + && messages[index].tool_call_id.is_some() + { + let tool = &messages[index]; + let tool_call_id = tool + .tool_call_id + .as_deref() + .expect("tool_call_id checked above"); + content.push(serde_json::json!({ "type": "tool_result", "tool_use_id": tool_call_id, - "content": convert_content_blocks(&message.content, false), - })] - } else if let Some(native) = native_anthropic_content(message) { - native - } else { - let mut blocks = convert_content_blocks(&message.content, message.role == "system"); - if let Some(tool_calls) = message - .tool_calls - .as_ref() - .filter(|calls| !calls.is_empty()) - { - for tool_call in tool_calls { - blocks.push(serde_json::json!({ - "type": "tool_use", - "id": tool_call.id, - "name": tool_call.name, - "input": tool_call.arguments, - })); - } + "content": convert_content_blocks(&tool.content, false), + })); + index += 1; + } + + // One turn may receive more than one steering message before the + // next model request. Keep their order while emitting one native + // Anthropic user message alongside the tool_result blocks. + while index < messages.len() && messages[index].role == "user" { + let steering = &messages[index]; + if let Some(native) = native_anthropic_content(steering) { + content.extend(native); + } else { + content.extend(convert_content_blocks(&steering.content, false)); } - blocks - }; - AnthropicMessage { role, content } - }) - .collect() + index += 1; + } + + converted.push(AnthropicMessage { + role: "user".to_string(), + content, + }); + continue; + } + + let role = message.role.clone(); + let content = if let Some(native) = native_anthropic_content(message) { + native + } else { + let mut blocks = convert_content_blocks(&message.content, message.role == "system"); + if let Some(tool_calls) = message + .tool_calls + .as_ref() + .filter(|calls| !calls.is_empty()) + { + for tool_call in tool_calls { + blocks.push(serde_json::json!({ + "type": "tool_use", + "id": tool_call.id, + "name": tool_call.name, + "input": tool_call.arguments, + })); + } + } + blocks + }; + converted.push(AnthropicMessage { role, content }); + index += 1; + } + + converted } fn native_anthropic_content(message: &Message) -> Option> { @@ -690,6 +729,35 @@ mod tests { assert_eq!(result["content"][1]["source"]["data"], "AAAA"); } + #[test] + fn tool_results_and_following_steering_share_one_user_content_array() { + let messages = vec![ + Message::tool("call_1", "lookup", "first result"), + Message::tool("call_2", "lookup", "second result"), + Message::user("用户补充指令"), + Message::user("再补充一条"), + Message::assistant("最终回答"), + ]; + + let converted = convert_messages(&messages); + + assert_eq!(converted.len(), 2); + assert_eq!(converted[0].role, "user"); + assert_eq!(converted[0].content.len(), 4); + assert_eq!(converted[0].content[0]["type"], "tool_result"); + assert_eq!(converted[0].content[0]["tool_use_id"], "call_1"); + assert_eq!(converted[0].content[1]["tool_use_id"], "call_2"); + assert_eq!( + converted[0].content[2], + json!({"type": "text", "text": "用户补充指令"}) + ); + assert_eq!( + converted[0].content[3], + json!({"type": "text", "text": "再补充一条"}) + ); + assert_eq!(converted[1].role, "assistant"); + } + #[test] fn native_stream_decodes_thinking_signature_tools_usage_and_replay_state() { let events = [ diff --git a/src/providers/openai.rs b/src/providers/openai.rs index 4fb0235..f542ec3 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -761,6 +761,46 @@ mod tests { assert_eq!(converted[1]["content"], "second image"); } + #[test] + fn assistant_tools_precede_all_tool_results_and_following_steering() { + let messages = vec![ + Message { + role: "assistant".to_string(), + content: vec![ContentBlock::text("calling tools")], + reasoning_content: None, + provider_state: None, + tool_call_id: None, + name: None, + tool_calls: Some(vec![ + ToolCall { + id: "call_1".to_string(), + name: "lookup".to_string(), + arguments: json!({"q": "one"}), + }, + ToolCall { + id: "call_2".to_string(), + name: "lookup".to_string(), + arguments: json!({"q": "two"}), + }, + ]), + }, + Message::tool("call_1", "lookup", "result"), + Message::tool("call_2", "lookup", "second result"), + Message::user("用户补充指令"), + ]; + + let converted = convert_messages(&messages); + + assert_eq!(converted.len(), 4); + assert_eq!(converted[0]["role"], "assistant"); + assert_eq!(converted[1]["role"], "tool"); + assert_eq!(converted[1]["tool_call_id"], "call_1"); + assert_eq!(converted[2]["role"], "tool"); + assert_eq!(converted[2]["tool_call_id"], "call_2"); + assert_eq!(converted[3]["role"], "user"); + assert_eq!(converted[3]["content"], "用户补充指令"); + } + #[test] fn assistant_images_are_never_serialized_as_native_content_parts() { let converted = convert_messages(&[Message { diff --git a/src/providers/stream.rs b/src/providers/stream.rs index c01da82..d161a5a 100644 --- a/src/providers/stream.rs +++ b/src/providers/stream.rs @@ -46,10 +46,7 @@ impl SseFramer { fn drain_frames(&mut self, finish: bool) -> Result, std::string::FromUtf8Error> { let mut frames = Vec::new(); - loop { - let Some((position, delimiter_len)) = find_sse_delimiter(&self.buffer) else { - break; - }; + while let Some((position, delimiter_len)) = find_sse_delimiter(&self.buffer) { let frame = self.buffer.drain(..position).collect::>(); self.buffer.drain(..delimiter_len); if let Some(data) = sse_data(frame)? { diff --git a/src/session/session.rs b/src/session/session.rs index 93cc109..79c037e 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -1,5 +1,5 @@ -use std::collections::HashMap; -use std::sync::Arc; +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex as StdMutex}; use tokio::sync::{Mutex, mpsc, oneshot}; @@ -198,7 +198,10 @@ pub enum HandleResult { } use crate::agent::context_compressor::ContextCompressionConfig; use crate::agent::system_prompt::build_system_prompt; -use crate::agent::{AgentError, AgentLoop, AgentTurnContext, ContextCompressor, TurnEmitter}; +use crate::agent::{ + AgentError, AgentLoop, AgentTurnContext, ContextCompressor, TurnEmitter, + steering::SteeringMailbox, +}; use crate::channels::slash_command::parse_slash_command; use crate::config::BrowserConfig; use crate::config::LLMProviderConfig; @@ -580,6 +583,10 @@ pub struct Session { /// Task queue for per-session serial agent processing agent_tx: Option>, + /// Monotonic admission order across steering and next-turn queue inputs. + /// It is allocated while holding the Session mutex so worker cleanup can + /// use that same lock as the send barrier. + next_task_sequence: u64, /// Cancel signal for the currently executing agent task current_cancel: Option>, active_turn_emitter: Option, @@ -603,19 +610,95 @@ pub struct Session { struct ActiveTurnEmitter { turn_id: String, emitter: TurnEmitter, + /// Mailbox for user messages that should be injected into this turn at a + /// safe agent-loop boundary. The mailbox is shared with AgentLoop via + /// `AgentTurnContext`; keeping it on the session handle makes admission + /// atomic with `/stop` and worker cleanup. + steering: Arc, + /// Original inbound tasks for accepted steering messages. ChatMessage + /// intentionally carries only durable history fields, so this side map + /// preserves channel context and rich MediaItem metadata if a terminal + /// turn has to fall back to a subsequent queued Turn. + recovery: StdArc>>, } /// A task to be processed by the per-session agent worker +#[derive(Clone)] struct AgentTask { channel: String, sender_id: String, chat_id: String, + sequence: u64, + client_message_id: Option, content: String, received_at: i64, media: Vec, channel_context: ChannelContext, } +/// Move terminally pending steering into the worker's local FIFO. The +/// recovery map is consulted first so the original channel context and rich +/// media metadata survive a same-turn fallback; only legacy/test producers +/// need the ChatMessage-derived fallback. +fn prepend_pending_steering( + mailbox: &SteeringMailbox, + recovery: &StdArc>>, + local_tasks: &mut VecDeque, + fallback_channel: &str, + fallback_chat_id: &str, + fallback_reply_to: &Option, + fallback_private: &HashMap, +) { + let pending = mailbox.close_and_take_pending(); + for message in pending.into_iter().rev() { + let recovered = recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&message.id); + let task = recovered.unwrap_or_else(|| { + let source = message.source.clone(); + AgentTask { + channel: source + .as_ref() + .and_then(|source| source.from_channel.clone()) + .unwrap_or_else(|| fallback_channel.to_string()), + sender_id: source + .as_ref() + .and_then(|source| source.from_user_id.clone()) + .unwrap_or_else(|| "unknown".to_string()), + chat_id: fallback_chat_id.to_string(), + sequence: 0, + client_message_id: Some(message.id.clone()), + content: message.content, + received_at: message.timestamp, + media: message + .media_refs + .into_iter() + .map(|media| MediaItem::new(media.path, media.media_type)) + .collect(), + channel_context: ChannelContext { + reply_to: fallback_reply_to.clone(), + private: fallback_private.clone(), + }, + } + }); + local_tasks.push_front(task); + } + recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); +} + +fn pop_lowest_sequence(local_tasks: &mut VecDeque) -> Option { + let index = local_tasks + .iter() + .enumerate() + .min_by_key(|(_, task)| task.sequence) + .map(|(index, _)| index)?; + local_tasks.remove(index) +} + #[derive(Clone)] struct AgentWorkerDeps { bus: Arc, @@ -678,6 +761,7 @@ impl Session { last_compressed_message_at: None, memory_manager, agent_tx: None, + next_task_sequence: 1, current_cancel: None, active_turn_emitter: None, worker_generation: 0, @@ -868,6 +952,7 @@ impl Session { last_compressed_message_at: restored_compressed_at, memory_manager, agent_tx: None, + next_task_sequence: 1, current_cancel: None, active_turn_emitter: None, worker_generation: 0, @@ -1009,6 +1094,12 @@ impl Session { .is_some_and(|active| active.turn_id == turn_id) } + fn allocate_task_sequence(&mut self) -> u64 { + let sequence = self.next_task_sequence; + self.next_task_sequence = self.next_task_sequence.wrapping_add(1).max(1); + sequence + } + #[cfg(test)] pub(super) fn state_version_for_test(&self) -> u64 { self.state_version @@ -1020,6 +1111,8 @@ impl Session { self.active_turn_emitter = Some(ActiveTurnEmitter { turn_id: turn_id.to_string(), emitter, + steering: SteeringMailbox::new_shared(), + recovery: StdArc::new(StdMutex::new(HashMap::new())), }); } @@ -1562,6 +1655,11 @@ pub static SLASH_COMMANDS: &[SlashCommand] = &[ description: "停止当前正在执行的任务并清空消息队列", aliases: &["/stop"], }, + SlashCommand { + name: "queue", + description: "等待当前任务完成后处理消息", + aliases: &["/queue"], + }, SlashCommand { name: "todo", description: "查看、完成或取消当前任务计划", @@ -1983,6 +2081,28 @@ impl SessionManager { let report = self.health.check().await; Ok((None, report.render_text())) } + "queue" => { + let sid = current_session_id + .ok_or_else(|| AgentError::Other("no active session".to_string()))?; + let content = args.map(str::trim).unwrap_or_default(); + if content.is_empty() { + return Err(AgentError::Other("Usage: /queue ".to_string())); + } + let session = self.get_or_create_session(sid).await?; + let task = AgentTask { + channel: channel.to_string(), + sender_id: "cli".to_string(), + chat_id: chat_id.to_string(), + sequence: 0, + client_message_id: None, + content: content.to_string(), + received_at: chrono::Utc::now().timestamp_millis(), + media: Vec::new(), + channel_context: ChannelContext::default(), + }; + self.enqueue_agent_task(session, sid.clone(), task).await?; + Ok((None, "消息已加入队列，将在当前任务完成后处理。".to_string())) + } "stop" => { let sid = current_session_id .ok_or_else(|| AgentError::Other("no active session".to_string()))?; @@ -1994,6 +2114,19 @@ impl SessionManager { msgs.push("当前任务已发送停止信号。".to_string()); } if let Some(active_turn) = guard.active_turn_emitter.take() { + // Closing the mailbox makes an in-flight ordinary + // input race resolve to the next queue (or an + // explicit queue-full response), rather than being + // accepted after `/stop` has invalidated the turn. + // `/stop` intentionally discards accepted-but-not-yet + // injected steering, matching its queue-clearing + // semantics. + let _ = active_turn.steering.close_and_take_pending(); + active_turn + .recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); active_turn.emitter.deactivate(); } if guard.agent_tx.take().is_some() { @@ -2691,6 +2824,37 @@ impl SessionManager { .scope(Some(unified_id.to_string()), async { // Check for slash command if let Some((cmd_name, cmd_args)) = parse_slash_command(content) { + // `/queue` is an input-routing primitive rather than a + // normal slash command. Preserve the original inbound + // metadata/media while stripping only the command + // prefix. Unlike steering, queued input is never + // admitted to the active turn. + if cmd_name == "queue" { + let queued_content = cmd_args.trim(); + if queued_content.is_empty() && inbound.media.is_empty() { + return Ok(HandleResult::CommandOutput( + "Usage: /queue ".to_string(), + )); + } + let task = AgentTask { + channel: channel.to_string(), + sender_id: sender_id.to_string(), + chat_id: chat_id.to_string(), + sequence: 0, + client_message_id: inbound + .client_message_id + .clone() + .filter(|id| !id.trim().is_empty()), + content: queued_content.to_string(), + received_at: inbound.received_at, + media: inbound.media.clone(), + channel_context: inbound.channel_context.clone(), + }; + return self + .enqueue_agent_task(session.clone(), unified_id.clone(), task) + .await; + } + let result = self .execute_slash_command( cmd_name, @@ -2714,83 +2878,159 @@ impl SessionManager { } // Normal message: enqueue to per-session worker for serial processing. - let task = AgentTask { + let mut task = AgentTask { channel: channel.to_string(), sender_id: sender_id.to_string(), chat_id: chat_id.to_string(), + sequence: 0, + client_message_id: inbound + .client_message_id + .clone() + .filter(|id| !id.trim().is_empty()), content: content.to_string(), received_at: inbound.received_at, media: inbound.media.clone(), channel_context: inbound.channel_context.clone(), }; - let session_clone = session.clone(); + + // While a turn is active, ordinary input is steering by + // default. The mailbox operation is synchronized with + // `/stop`: if it is full or already closed, fall back to the + // bounded session queue instead of dropping the input. let unified_str = unified_id.to_string(); - { - let mut guard = session_clone.lock().await; - let needs_spawn = guard.agent_tx.is_none() - || guard.agent_tx.as_ref().is_some_and(|tx| tx.is_closed()); - if needs_spawn { - guard.agent_tx = None; - guard.current_cancel = None; - guard.worker_generation = guard.worker_generation.wrapping_add(1); - let generation = guard.worker_generation; - let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); - guard.agent_tx = Some(tx); - spawn_agent_worker( - rx, - session_clone.clone(), - self.worker_deps(), - generation, - unified_str.clone(), - ); - } - let Some(agent_tx) = guard.agent_tx.as_ref() else { - return Err(AgentError::Other( - "agent worker queue was not initialized".to_string(), - )); + let mut guard = session.lock().await; + if task.sequence == 0 { + task.sequence = guard.allocate_task_sequence(); + } + if let Some(active) = guard.active_turn_emitter.as_ref() { + let media_refs: Vec = + inbound.media.iter().map(MediaItem::to_media_ref).collect(); + let source = MessageSource { + kind: SourceKind::UserInput, + from_channel: Some(channel.to_string()), + from_session: None, + from_user_id: Some(sender_id.to_string()), + system_name: None, + task_id: None, }; - if let Err(e) = agent_tx.try_send(task) { - if matches!(e, mpsc::error::TrySendError::Full(_)) { - tracing::warn!(session_id = %unified_str, capacity = SESSION_QUEUE_CAPACITY, "Session queue is full"); - return Ok(HandleResult::CommandOutput( - "当前对话消息队列已满，请稍后重试。".to_string(), - )); + let mut message = + guard.create_user_message_with_source(content, media_refs, source); + if let Some(id) = inbound + .client_message_id + .as_deref() + .filter(|id| !id.trim().is_empty()) + { + message.id = id.to_string(); + } + message.timestamp = inbound.received_at; + let message_id = message.id.clone(); + let mut recovery_task = task.clone(); + recovery_task.client_message_id = Some(message_id.clone()); + active + .recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert(message_id.clone(), recovery_task); + match active.steering.try_push(message) { + Ok(()) => return Ok(HandleResult::AgentProcessing), + Err(_) => { + active + .recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&message_id); } - // Worker died after we just spawned it — respawn with the recovered task - let task = e.into_inner(); - guard.agent_tx = None; - guard.current_cancel = None; - guard.worker_generation = guard.worker_generation.wrapping_add(1); - let generation = guard.worker_generation; - let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); - guard.agent_tx = Some(tx); - spawn_agent_worker( - rx, - session_clone.clone(), - self.worker_deps(), - generation, - unified_str.clone(), - ); - guard - .agent_tx - .as_ref() - .ok_or_else(|| { - AgentError::Other( - "agent worker queue was not initialized".to_string(), - ) - })? - .try_send(task) - .map_err(|_| { - AgentError::Other( - "agent worker spawn+send failed irrecoverably".to_string(), - ) - })?; } } - Ok(HandleResult::AgentProcessing) + + self.enqueue_agent_task_locked(&mut guard, &session, &unified_str, task) }) .await } + + /// Enqueue one user task on the session's bounded serial worker. This is + /// shared by ordinary messages, `/queue`, and the reliable fallback when + /// steering admission races with mailbox closure/full capacity. + async fn enqueue_agent_task( + &self, + session: Arc>, + unified_id: UnifiedSessionId, + task: AgentTask, + ) -> Result { + let unified_str = unified_id.to_string(); + let mut guard = session.lock().await; + self.enqueue_agent_task_locked(&mut guard, &session, &unified_str, task) + } + + fn enqueue_agent_task_locked( + &self, + guard: &mut Session, + session: &Arc>, + unified_str: &str, + mut task: AgentTask, + ) -> Result { + if task.sequence == 0 { + task.sequence = guard.allocate_task_sequence(); + } + let needs_spawn = + guard.agent_tx.is_none() || guard.agent_tx.as_ref().is_some_and(|tx| tx.is_closed()); + if needs_spawn { + guard.agent_tx = None; + guard.current_cancel = None; + guard.worker_generation = guard.worker_generation.wrapping_add(1); + let generation = guard.worker_generation; + let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); + guard.agent_tx = Some(tx); + spawn_agent_worker( + rx, + session.clone(), + self.worker_deps(), + generation, + unified_str.to_string(), + ); + } + let Some(agent_tx) = guard.agent_tx.as_ref() else { + return Err(AgentError::Other( + "agent worker queue was not initialized".to_string(), + )); + }; + if let Err(error) = agent_tx.try_send(task) { + if matches!(error, mpsc::error::TrySendError::Full(_)) { + tracing::warn!(session_id = %unified_str, capacity = SESSION_QUEUE_CAPACITY, "Session queue is full"); + return Ok(HandleResult::CommandOutput( + "当前对话消息队列已满，请稍后重试。".to_string(), + )); + } + + // The worker can exit between spawn and send. Recreate the + // worker under the same lock and retry the recovered task once. + let task = error.into_inner(); + guard.agent_tx = None; + guard.current_cancel = None; + guard.worker_generation = guard.worker_generation.wrapping_add(1); + let generation = guard.worker_generation; + let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); + guard.agent_tx = Some(tx); + spawn_agent_worker( + rx, + session.clone(), + self.worker_deps(), + generation, + unified_str.to_string(), + ); + guard + .agent_tx + .as_ref() + .ok_or_else(|| { + AgentError::Other("agent worker queue was not initialized".to_string()) + })? + .try_send(task) + .map_err(|_| { + AgentError::Other("agent worker spawn+send failed irrecoverably".to_string()) + })?; + } + Ok(HandleResult::AgentProcessing) + } } async fn generate_title( @@ -2886,7 +3126,28 @@ fn spawn_agent_worker( task_supervisor.spawn(format!("session-worker:{unified_str}"), async move { let unified_for_source = unified_str.clone(); let _scope = CURRENT_SOURCE_SESSION.scope(Some(unified_for_source), async { - 'tasks: while let Some(task) = task_rx.recv().await { + // Steering that cannot be consumed by the current turn is + // prepended here. This local deque is deliberately ahead of the + // channel queue: messages admitted after the mailbox closed must + // not overtake an earlier accepted steering message, and moving + // pending inputs here cannot fail due to channel capacity. + let mut local_tasks = VecDeque::new(); + 'tasks: loop { + // Admission sequence numbers are allocated while holding the + // Session lock. Drain everything currently visible on the + // channel before selecting the smallest sequence, so a + // terminal fallback cannot overtake an earlier `/queue` task. + while let Ok(task) = task_rx.try_recv() { + local_tasks.push_back(task); + } + let task = if let Some(task) = pop_lowest_sequence(&mut local_tasks) { + task + } else { + match task_rx.recv().await { + Some(task) => task, + None => break, + } + }; let task_chan = task.channel.clone(); let task_cid = task.chat_id.clone(); let task_metadata = task.channel_context.private.clone(); @@ -2913,6 +3174,13 @@ fn spawn_agent_worker( }; let mut message = guard.create_user_message_with_source(&task.content, media_refs, source); + if let Some(id) = task + .client_message_id + .as_deref() + .filter(|id| !id.trim().is_empty()) + { + message.id = id.to_string(); + } message.timestamp = task.received_at; message }; @@ -2938,6 +3206,12 @@ fn spawn_agent_worker( system_prompt_out, base_version, cancel_rx, + turn_controller, + turn_emitter, + turn_receiver, + initial_turn, + steering, + recovery, ) = { let mut guard = session.lock().await; @@ -2973,6 +3247,26 @@ fn spawn_agent_worker( } guard.current_cancel = Some(cancel_tx); + // Install the active-turn handle before memory recall and + // context preparation. This closes the historical race + // where a message received during slow preparation was + // queued as a new turn instead of steering the running + // one. + let (turn_controller, turn_emitter, turn_receiver) = + TurnController::start( + unified_str.clone(), + uuid::Uuid::new_v4().to_string(), + ); + let initial_turn = turn_controller.snapshot(); + let steering = SteeringMailbox::new_shared(); + let recovery = StdArc::new(StdMutex::new(HashMap::new())); + guard.active_turn_emitter = Some(ActiveTurnEmitter { + turn_id: initial_turn.id.0.clone(), + emitter: turn_emitter.clone(), + steering: steering.clone(), + recovery: recovery.clone(), + }); + ( agent, history_raw, @@ -2980,6 +3274,12 @@ fn spawn_agent_worker( guard.build_system_prompt(&skills_prompt), guard.state_version, cancel_rx, + turn_controller, + turn_emitter, + turn_receiver, + initial_turn, + steering, + recovery, ) }; // lock released @@ -2996,6 +3296,29 @@ fn spawn_agent_worker( let meta_snapshot = { let mut guard = session.lock().await; if guard.worker_generation != worker_gen { + steering.close_and_take_pending(); + recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); + if guard + .active_turn_emitter + .as_ref() + .is_some_and(|active| active.turn_id == initial_turn.id.0) + && let Some(active) = guard.active_turn_emitter.take() + { + let _ = active.steering.close_and_take_pending(); + active + .recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); + active.emitter.deactivate(); + } + turn_emitter.deactivate(); + turn_controller.cancel(Some( + "session changed before turn preparation completed".to_string(), + )); return; } if guard.state_version != base_version { @@ -3004,6 +3327,27 @@ fn spawn_agent_worker( "Session changed while preparing agent history; dropping stale task" ); guard.current_cancel = None; + if guard + .active_turn_emitter + .as_ref() + .is_some_and(|active| active.turn_id == initial_turn.id.0) + && let Some(active) = guard.active_turn_emitter.take() + { + prepend_pending_steering( + &active.steering, + &active.recovery, + &mut local_tasks, + &task_chan, + &task_cid, + &task_reply_to, + &task_metadata, + ); + active.emitter.deactivate(); + } + turn_emitter.deactivate(); + turn_controller.cancel(Some( + "session changed while preparing agent history".to_string(), + )); continue 'tasks; } if prepared_input.created_timelines { @@ -3021,11 +3365,6 @@ fn spawn_agent_worker( let history_out = prepared_input.messages; let runtime_context = prepared_input.runtime; - let (turn_controller, turn_emitter, turn_receiver) = TurnController::start( - unified_str.clone(), - uuid::Uuid::new_v4().to_string(), - ); - let initial_turn = turn_controller.snapshot(); let active_turn_id = initial_turn.id.0.clone(); let turn_target = crate::channels::TurnTarget { channel: task_chan.clone(), @@ -3052,6 +3391,23 @@ fn spawn_agent_worker( { let mut guard = session.lock().await; if guard.worker_generation != worker_gen || guard.state_version != base_version { + if guard + .active_turn_emitter + .as_ref() + .is_some_and(|active| active.turn_id == active_turn_id) + && let Some(active) = guard.active_turn_emitter.take() + { + prepend_pending_steering( + &active.steering, + &active.recovery, + &mut local_tasks, + &task_chan, + &task_cid, + &task_reply_to, + &task_metadata, + ); + active.emitter.deactivate(); + } turn_emitter.deactivate(); turn_controller.cancel(Some( "session changed before model execution".to_string(), @@ -3059,15 +3415,12 @@ fn spawn_agent_worker( guard.current_cancel = None; continue 'tasks; } - guard.active_turn_emitter = Some(ActiveTurnEmitter { - turn_id: initial_turn.id.0.clone(), - emitter: turn_emitter.clone(), - }); } - let agent_turn = AgentTurnContext::new( + let agent_turn = AgentTurnContext::new_with_steering( initial_turn.id.0.clone(), initial_turn.message_id.clone(), turn_emitter, + steering.clone(), ); // Phase 2 + 3: LLM call with cancellation @@ -3084,6 +3437,7 @@ fn spawn_agent_worker( let turn_lifecycle = &turn_controller; let pending_turn_deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); let scoped_turn_deliveries = pending_turn_deliveries.clone(); + let steering_for_process = steering.clone(); let process_future = async move { let response_session_id = unified_str2.clone(); let tool_context = ToolExecutionContext::for_session(&response_session_id) @@ -3258,9 +3612,23 @@ fn spawn_agent_worker( .map(|value| value.prompt_tokens); let (provider_name, model_name) = { let guard = session2.lock().await; - if guard.worker_generation != worker_gen - || guard.state_version != base_version - { + if guard.worker_generation != worker_gen { + // A generation change is the explicit `/stop` / + // worker-replacement boundary; discard pending + // and in-flight steering just like the command + // path does. + steering_for_process.close_and_take_pending(); + turn_lifecycle.cancel(Some( + "session changed before turn commit".to_string(), + )); + return; + } + if guard.state_version != base_version { + // A version-only mutation (for example /clear) + // invalidates this commit but does not carry the + // explicit stop/discard semantics. Requeue the + // accepted steering for the next local Turn. + steering_for_process.restore_drained(); turn_lifecycle.cancel(Some( "session changed before turn commit".to_string(), )); @@ -3303,6 +3671,9 @@ fn spawn_agent_worker( .await { Ok(committed_messages) => { + // AgentLoop keeps drained steering in the + // mailbox until the whole Turn batch is durable. + steering_for_process.commit_drained(); let mut guard = session2.lock().await; let prompt_message_count = guard.messages.len().saturating_sub(1); guard @@ -3312,6 +3683,10 @@ fn spawn_agent_worker( } Err(e) => { tracing::error!(error = %e, "Failed to atomically persist agent turn"); + // Restore the in-memory steering batch so an + // accepted user input is retried as a subsequent + // Turn instead of silently disappearing. + steering_for_process.restore_drained(); None } }; @@ -3407,6 +3782,22 @@ fn spawn_agent_worker( .await; } + // AgentLoop closes its mailbox on terminal/error and leaves + // any not-yet-injected messages pending. Transfer those + // accepted inputs to the ordinary bounded queue so they are + // processed by a subsequent Turn instead of being lost. + // `/stop` calls `close_and_take_pending` itself, therefore + // its intentionally discarded inputs do not reach here. + prepend_pending_steering( + &steering, + &recovery, + &mut local_tasks, + &task_chan, + &task_cid, + &task_reply_to, + &task_metadata, + ); + // Clean up let mut guard = session.lock().await; if guard @@ -3415,10 +3806,25 @@ fn spawn_agent_worker( .is_some_and(|active| active.turn_id == active_turn_id) && let Some(active) = guard.active_turn_emitter.take() { + active.steering.close(); + active + .recovery + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); active.emitter.deactivate(); } if guard.worker_generation == worker_gen { - guard.current_cancel = None; + if local_tasks.is_empty() { + guard.current_cancel = None; + } else { + // Keep the session observable as busy while local + // fallback tasks are waiting for their next Turn. + // `/stop` can still invalidate this generation before + // the worker starts the next task. + let (cancel_tx, _cancel_rx) = oneshot::channel(); + guard.current_cancel = Some(cancel_tx); + } } } }).await; @@ -3582,7 +3988,14 @@ fn format_task_notification( #[cfg(test)] mod slash_command_tests { - use super::resolve_slash_command; + use super::{ + AgentTask, SLASH_COMMANDS, pop_lowest_sequence, prepend_pending_steering, + resolve_slash_command, + }; + use crate::agent::steering::{SteeringMailbox, SteeringPushError}; + use crate::bus::{ChannelContext, ChatMessage, MediaItem}; + use std::collections::{HashMap, VecDeque}; + use std::sync::{Arc, Mutex}; #[test] fn aliases_resolve_to_their_canonical_command() { @@ -3608,4 +4021,136 @@ mod slash_command_tests { ); assert!(resolve_slash_command("unknown").is_none()); } + + #[test] + fn queue_is_advertised_and_resolves_to_the_routing_command() { + assert!(SLASH_COMMANDS.iter().any(|command| command.name == "queue")); + assert_eq!( + resolve_slash_command("/queue").map(|command| command.name), + Some("queue") + ); + } + + #[test] + fn steering_capacity_error_is_reliable_fallback_signal() { + let mailbox = SteeringMailbox::with_limits(1, usize::MAX); + mailbox.try_push(ChatMessage::user("first")).unwrap(); + let error = mailbox.try_push(ChatMessage::user("second")).unwrap_err(); + assert!(error.is_full()); + assert!(matches!(error, SteeringPushError::Full(_))); + assert_eq!(mailbox.len(), 1); + } + + #[test] + fn terminal_pending_steering_keeps_original_route_and_media() { + let mailbox = SteeringMailbox::new_shared(); + let recovery = Arc::new(Mutex::new(HashMap::new())); + let mut message = ChatMessage::user("inspect the upload"); + message.id = "client-steer-1".to_string(); + message.media_refs.push(crate::bus::MediaRef { + path: "/tmp/upload.bin".to_string(), + media_type: "file".to_string(), + }); + let original = AgentTask { + channel: "feishu".to_string(), + sender_id: "user-7".to_string(), + chat_id: "chat-9".to_string(), + sequence: 7, + client_message_id: Some(message.id.clone()), + content: message.content.clone(), + received_at: message.timestamp, + media: vec![MediaItem::new("/tmp/upload.bin", "file")], + channel_context: ChannelContext { + reply_to: Some("reply-1".to_string()), + private: HashMap::from([(String::from("thread"), String::from("root-1"))]), + }, + }; + recovery + .lock() + .unwrap() + .insert(message.id.clone(), original); + mailbox.try_push(message).unwrap(); + + let mut local_tasks = VecDeque::new(); + prepend_pending_steering( + &mailbox, + &recovery, + &mut local_tasks, + "cli", + "fallback-chat", + &None, + &HashMap::new(), + ); + let recovered = local_tasks.pop_front().unwrap(); + assert_eq!(recovered.channel, "feishu"); + assert_eq!(recovered.sender_id, "user-7"); + assert_eq!(recovered.chat_id, "chat-9"); + assert_eq!( + recovered.channel_context.reply_to.as_deref(), + Some("reply-1") + ); + assert_eq!( + recovered.channel_context.private.get("thread").unwrap(), + "root-1" + ); + assert_eq!(recovered.media[0].path, "/tmp/upload.bin"); + assert!(mailbox.is_closed()); + assert!(recovery.lock().unwrap().is_empty()); + } + + #[test] + fn pending_and_channel_tasks_are_selected_by_admission_sequence() { + let mailbox = SteeringMailbox::new_shared(); + let recovery = Arc::new(Mutex::new(HashMap::new())); + let mut pending = ChatMessage::user("pending"); + pending.id = "pending-id".to_string(); + let template = AgentTask { + channel: "cli".to_string(), + sender_id: "user".to_string(), + chat_id: "chat".to_string(), + sequence: 10, + client_message_id: Some(pending.id.clone()), + content: pending.content.clone(), + received_at: pending.timestamp, + media: Vec::new(), + channel_context: ChannelContext::default(), + }; + recovery + .lock() + .unwrap() + .insert(pending.id.clone(), template.clone()); + mailbox.try_push(pending).unwrap(); + + let mut local_tasks = VecDeque::from([ + AgentTask { + sequence: 30, + ..template.clone() + }, + AgentTask { + sequence: 20, + ..template + }, + ]); + prepend_pending_steering( + &mailbox, + &recovery, + &mut local_tasks, + "cli", + "chat", + &None, + &HashMap::new(), + ); + assert_eq!(pop_lowest_sequence(&mut local_tasks).unwrap().sequence, 10); + assert_eq!(pop_lowest_sequence(&mut local_tasks).unwrap().sequence, 20); + assert_eq!(pop_lowest_sequence(&mut local_tasks).unwrap().sequence, 30); + } + + #[test] + fn stop_style_close_discards_pending_steering() { + let mailbox = SteeringMailbox::new_shared(); + mailbox.try_push(ChatMessage::user("discard me")).unwrap(); + assert_eq!(mailbox.close_and_take_pending().len(), 1); + assert!(mailbox.is_closed()); + assert!(mailbox.take_pending().is_empty()); + } } diff --git a/src/tools/pty.rs b/src/tools/pty.rs index e4b1951..22e793b 100644 --- a/src/tools/pty.rs +++ b/src/tools/pty.rs @@ -307,7 +307,7 @@ impl PtyManager { let total = guard.output_total_lines; let buffer_len = guard.output_buffer.len(); - let start = 0_usize.max(offset); + let start = offset; let skip_old = total.saturating_sub(buffer_len); let view_start = start.saturating_sub(skip_old); diff --git a/tests/test_request_format.rs b/tests/test_request_format.rs index 0deff88..9d9394f 100644 --- a/tests/test_request_format.rs +++ b/tests/test_request_format.rs @@ -167,6 +167,7 @@ fn test_user_input_accepts_upload_ids_and_old_payloads() { let message = WsInbound::UserInput { content: "处理文件".to_string(), upload_ids: vec!["upload-1".to_string()], + client_message_id: None, channel: None, chat_id: None, sender_id: None, diff --git a/webui/package-lock.json b/webui/package-lock.json index 21cb297..80b2336 100644 --- a/webui/package-lock.json +++ b/webui/package-lock.json @@ -1,12 +1,12 @@ { "name": "picobot-webui", - "version": "1.4.1", + "version": "1.5.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "picobot-webui", - "version": "1.4.1", + "version": "1.5.0", "dependencies": { "bits-ui": "^2.0.0", "dompurify": "^3.4.12", diff --git a/webui/package.json b/webui/package.json index 8084cdd..a75e747 100644 --- a/webui/package.json +++ b/webui/package.json @@ -1,7 +1,7 @@ { "name": "picobot-webui", "private": true, - "version": "1.4.1", + "version": "1.5.0", "type": "module", "engines": { "node": ">=20" diff --git a/webui/src/pages/ChatPage.svelte b/webui/src/pages/ChatPage.svelte index 1a13ba7..8b98229 100644 --- a/webui/src/pages/ChatPage.svelte +++ b/webui/src/pages/ChatPage.svelte @@ -152,9 +152,22 @@ } case "turn_committed": { if (frame.session_id !== currentId || frame.history_revision <= historyRevision) break; - const byId = new Map(messages.map((message) => [message.id, message])); - for (const message of frame.messages || []) byId.set(message.id, message); - messages = [...byId.values()]; + const committed = frame.messages || []; + const committedIds = new Set(committed.map((message) => message.id)); + const firstOptimisticMatch = messages.findIndex((message) => committedIds.has(message.id)); + const retained = messages.filter((message) => !committedIds.has(message.id)); + const insertionIndex = firstOptimisticMatch < 0 + ? retained.length + : messages.slice(0, firstOptimisticMatch) + .filter((message) => !committedIds.has(message.id)).length; + // Replace optimistic steering at its original position with the whole + // durable batch. A Map#set replacement would keep only the user's + // old slot and incorrectly place earlier tool messages after it. + messages = [ + ...retained.slice(0, insertionIndex), + ...committed, + ...retained.slice(insertionIndex) + ]; historyRevision = frame.history_revision; if (activeTurn?.status !== "running" && (frame.messages || []).some((message) => message.id === activeTurn?.message_id)) { @@ -216,15 +229,21 @@ const content = draft.trim(); const ready = pendingUploads.filter((upload) => upload.status === "ready"); if ((!content && !ready.length) || !chat.connected || pendingUploads.some((upload) => upload.status === "uploading")) return; + const clientMessageId = randomId(); appendMessage("user", content, ready.map((upload, index) => ({ index, name: upload.name, media_type: upload.media_type, mime_type: upload.mime_type, local_url: upload.localUrl - }))); + })), clientMessageId); thinking = true; - send({ type: "user_input", content, upload_ids: ready.map((upload) => upload.upload_id) }); + send({ + type: "user_input", + content, + upload_ids: ready.map((upload) => upload.upload_id), + client_message_id: clientMessageId + }); draft = ""; pendingUploads = []; commandMenuDismissed = false;