mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
V1
This commit is contained in:
@@ -0,0 +1,694 @@
|
||||
## Multi‑Agent Collaboration Design for `codex-core`
|
||||
|
||||
This document sketches how to layer a multi‑agent collaboration system on top of the existing `codex-core` agent loop (sessions, turns, tools, context manager), using the proposed `collaboration.*` tools.
|
||||
|
||||
The focus is on:
|
||||
|
||||
- Keeping the current single‑agent semantics intact.
|
||||
- Treating “agents” as *logical participants* within a single Codex session.
|
||||
- Reusing existing primitives (`Session`, `TurnContext`, `ContextManager`, `ToolsConfig`, tasks) wherever possible.
|
||||
- Providing clear limits (max agents, max depth, per‑agent “juice”).
|
||||
|
||||
---
|
||||
|
||||
## 1. Mental Model
|
||||
|
||||
- A **Codex session** remains the top‑level unit, with one running `SessionTask` (typically `RegularTask`) at a time.
|
||||
- Within a session we introduce a **Collaboration Graph**: a set of logical agents that can exchange messages and call tools, orchestrated entirely within the same session.
|
||||
- The **main agent** corresponds to the existing single agent Codex exposes today. It gets the user’s prompt and controls collaboration via the new `collaboration.*` tools.
|
||||
- **Child agents** are lightweight logical agents:
|
||||
- Each has its own conversation history (`ContextManager`), view of tools, and model configuration.
|
||||
- Each runs its own *logical agentic loop* in a separate Tokio task, continuously consuming its model stream so that server‑side buffers do not depend on the parent calling `wait`.
|
||||
- All agents share the same session environment (cwd, sandbox manager, MCP connections, rollout recorder), but their prompts and tool calls are built from their own per‑agent `ContextManager`.
|
||||
- `collaboration.wait` becomes a **coordination primitive**: the main agent uses it to block until child agents have made enough progress (tokens/messages) rather than to drive the child loops directly.
|
||||
|
||||
Agents, from Codex’s point of view, are essentially:
|
||||
|
||||
```rust
|
||||
struct AgentId(pub u32); // 0-based index
|
||||
|
||||
struct AgentState {
|
||||
id: AgentId,
|
||||
parent: Option<AgentId>,
|
||||
depth: u32,
|
||||
|
||||
/// Configuration used to build per-agent TurnContexts.
|
||||
config: SessionConfiguration,
|
||||
|
||||
/// Agent-local conversation history (what this agent “remembers”).
|
||||
history: ContextManager,
|
||||
|
||||
/// Human-readable instructions / persona for this agent.
|
||||
instructions: Option<String>,
|
||||
|
||||
/// Soft “compute budget” in tokens for this agent.
|
||||
max_juice: Option<u32>,
|
||||
used_juice: u32,
|
||||
|
||||
/// High-level lifecycle state.
|
||||
status: AgentLifecycleState,
|
||||
}
|
||||
```
|
||||
|
||||
The main agent is `AgentId(0)` and is backed by the existing session‑wide `SessionConfiguration` and `SessionState.history`.
|
||||
|
||||
Agent lifecycle is modeled explicitly as:
|
||||
|
||||
```rust
|
||||
enum AgentLifecycleState {
|
||||
/// Agent is not currently streaming; the last assistant message
|
||||
/// (if any) is captured for easy retrieval.
|
||||
Idle { last_agent_message: String },
|
||||
/// Agent is in the middle of a turn / streaming response.
|
||||
Running,
|
||||
/// Agent hit a fatal error; `error` contains a concise summary.
|
||||
Error { error: String },
|
||||
/// Agent was explicitly closed (via `collaboration.close`); no
|
||||
/// further messages or tool calls are allowed.
|
||||
Closed,
|
||||
/// Agent is blocked waiting on a user approval (exec / patch /
|
||||
/// other), summarised as `request`.
|
||||
WaitingForApproval { request: String },
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Collaboration State and Limits
|
||||
|
||||
We introduce a session‑scoped **CollaborationState** owned by `Session`:
|
||||
|
||||
```rust
|
||||
struct CollaborationLimits {
|
||||
max_agents: u32,
|
||||
max_depth: u32,
|
||||
}
|
||||
|
||||
struct CollaborationState {
|
||||
/// Ordered list of agents; index is the AgentId.
|
||||
agents: Vec<AgentState>,
|
||||
/// Invariants and safety limits.
|
||||
limits: CollaborationLimits,
|
||||
/// Adjacency for the agent tree: parent → children.
|
||||
children: HashMap<AgentId, Vec<AgentId>>,
|
||||
}
|
||||
```
|
||||
|
||||
Key properties:
|
||||
|
||||
- **Lifetime**: `CollaborationState` lives for the duration of the session and persists across turns.
|
||||
- **Root agent**:
|
||||
- Created implicitly when collaboration is first used.
|
||||
- Mirrors the current `SessionConfiguration` and `SessionState.history`.
|
||||
- Has `parent = None`, `depth = 0`.
|
||||
- **Limits**:
|
||||
- `max_agents`: total agents in `agents` (including root); enforced on `collaboration.init_agent`.
|
||||
- `max_depth`: maximum `AgentState.depth`; child creation fails if `parent.depth + 1 > max_depth`.
|
||||
- **Juice**:
|
||||
- Each agent has `max_juice` (optional) and `used_juice`.
|
||||
- `used_juice` is derived from per‑agent `ContextManager` token usage (`get_total_token_usage`) and/or incremental deltas per turn.
|
||||
- **Tree topology**:
|
||||
- Agents form a strict **tree**:
|
||||
- Each `AgentState` has exactly one `parent` (except the root) and zero or more children.
|
||||
- A parent can only interact with its **direct children** (depth‑1 relative to the parent), never grandchildren or siblings.
|
||||
- All collaboration tools enforce this:
|
||||
- `send`, `wait`, `close` (see §5.5) only accept agents that are direct children of the caller.
|
||||
- This keeps reasoning local and prevents tangled cross‑branch coordination.
|
||||
|
||||
Storage location:
|
||||
|
||||
- Conceptually, `CollaborationState` is part of the session’s mutable state.
|
||||
- It can be stored alongside `SessionState` (e.g., as an optional field) or in a sibling struct owned by `Session` and protected by a mutex:
|
||||
- `SessionState` continues to own the *user‑facing* history.
|
||||
- `CollaborationState` owns per‑agent histories, which may be a subset or re‑projection of that history.
|
||||
|
||||
---
|
||||
|
||||
## 3. Agent Context vs Session Context
|
||||
|
||||
### 3.1 Session Configuration and TurnContext
|
||||
|
||||
`Session` already owns:
|
||||
|
||||
- `SessionConfiguration` (model, cwd, sandbox/approval policies, features, etc.).
|
||||
- `SessionState.history: ContextManager` (session‑wide conversation transcript).
|
||||
- `TurnContext` for each user turn (per‑turn model client + environment).
|
||||
|
||||
We extend this with **per‑agent configuration**:
|
||||
|
||||
- Each `AgentState.config` is a `SessionConfiguration`:
|
||||
- Initially cloned from the current session configuration.
|
||||
- Overridden by fields provided in `CollaborationInitAgentInput` (model, sandbox policy, approval policy, etc.).
|
||||
- When Codex wants to run that agent, it builds a per‑agent `TurnContext` by reusing `Session::make_turn_context`:
|
||||
|
||||
```rust
|
||||
fn make_agent_turn_context(
|
||||
session: &Session,
|
||||
agent: &AgentState,
|
||||
sub_id: String,
|
||||
) -> TurnContext {
|
||||
Session::make_turn_context(
|
||||
Some(Arc::clone(&session.services.auth_manager)),
|
||||
&session.services.otel_event_manager,
|
||||
agent.config.provider.clone(),
|
||||
&agent.config,
|
||||
session.conversation_id,
|
||||
sub_id,
|
||||
)
|
||||
}
|
||||
```
|
||||
|
||||
This keeps all the existing plumbing (telemetry, model selection, truncation policy, sandboxing) intact.
|
||||
|
||||
### 3.2 Agent Histories vs Session History
|
||||
|
||||
We distinguish:
|
||||
|
||||
- **Session/global history** (`SessionState.history`): what the user sees in the UI and what the main agent uses as its context. This remains the canonical transcript of the conversation.
|
||||
- **Agent history** (`AgentState.history`): what each logical agent sees when it reasons:
|
||||
- Includes messages *to and from* that agent, as well as tool calls initiated by that agent.
|
||||
- Represented as a `ContextManager`, so we retain all invariants and token accounting (call/output pairing, compaction, etc.).
|
||||
|
||||
Synchronization strategy:
|
||||
|
||||
- The main agent’s history is exactly `SessionState.history`.
|
||||
- For child agents:
|
||||
- On creation, we initialize `AgentState.history` according to `ContextStrategy` (see below).
|
||||
- When an agent sends a message or triggers a tool:
|
||||
- We append the corresponding `ResponseItem`s into that agent’s `AgentState.history`.
|
||||
- We also emit a *summarized* representation into `SessionState.history` (e.g., a compact `Message` describing “Agent 2 replied with: …”) so the user can follow the collaboration.
|
||||
- This keeps per‑agent context accurate while avoiding flooding the main transcript with every internal detail.
|
||||
|
||||
---
|
||||
|
||||
## 4. `ContextStrategy` Semantics
|
||||
|
||||
The proposed enum:
|
||||
|
||||
```rust
|
||||
enum ContextStrategy {
|
||||
#[default]
|
||||
New,
|
||||
Fork,
|
||||
Replace(Vec<ResponseItem>),
|
||||
}
|
||||
```
|
||||
|
||||
applies when creating a new agent:
|
||||
|
||||
- `New`:
|
||||
- Agent starts with an *empty* logical history, except for any synthetic system/role messages Codex inserts (see instructions below).
|
||||
- Useful for “specialist” agents that should not see the full prior conversation.
|
||||
- `Fork`:
|
||||
- The new agent’s `ContextManager` is initialized by cloning the parent agent’s promptable history:
|
||||
- `parent.history.get_history_for_prompt()` is a good starting point.
|
||||
- This gives the child agent the same conversational context as the parent at the time of the fork.
|
||||
- `Replace(Vec<ResponseItem>)`:
|
||||
- A more advanced mode, where the caller provides the full history for the agent.
|
||||
- In v1 we can limit usage to internal tooling / tests, but keeping it in the API makes training and “rehydration from logs” possible later.
|
||||
|
||||
In all cases, we also inject a small system message at the top of the agent’s history describing:
|
||||
|
||||
- The agent’s id and role.
|
||||
- Any initial instructions.
|
||||
- Its constraints (depth, max_juice, allowed tools).
|
||||
|
||||
This system message is created as a `ResponseItem::Message { role: "system", content: … }` and only lives in `AgentState.history`, not in the global transcript.
|
||||
|
||||
---
|
||||
|
||||
## 5. Tooling: `collaboration.*` Tools
|
||||
|
||||
We introduce five tools:
|
||||
|
||||
- `collaboration.init_agent`
|
||||
- `collaboration.send`
|
||||
- `collaboration.wait`
|
||||
- `collaboration.get_state`
|
||||
- `collaboration.close`
|
||||
|
||||
They are exposed as standard function tools via `ToolSpec::Function` and implemented as `ToolHandler`s that operate on `Session` + `TurnContext` and `CollaborationState`.
|
||||
|
||||
All tools share:
|
||||
|
||||
- A common metadata bag:
|
||||
|
||||
```rust
|
||||
pub struct ExtraMetadata(pub HashMap<String, serde_json::Value>);
|
||||
```
|
||||
|
||||
- A shared convention:
|
||||
- `message_tool_call_success`: true iff the tool handled the request successfully (even if no agent ran).
|
||||
- `message_tool_call_error_should_penalize_model`: true when the error is clearly model‑caused (e.g., invalid agent index), false for environment/limit issues.
|
||||
- `extra`: structured fields (agent ids, state snapshots, counters) that can evolve without changing the top‑level schema.
|
||||
|
||||
### 5.1 `collaboration.init_agent`
|
||||
|
||||
**Purpose**
|
||||
|
||||
Create a new logical agent as a child of the calling agent.
|
||||
|
||||
**Input**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationInitAgentInput {
|
||||
/// Index proposed by the model; ignored by the server.
|
||||
agent_idx: u32,
|
||||
|
||||
/// New fields (all default to current agent values).
|
||||
context_strategy: Option<ContextStrategy>,
|
||||
instructions: Option<String>,
|
||||
sandbox_policy: Option<SandboxPolicy>,
|
||||
|
||||
/// Future: per-agent model choice.
|
||||
model: String,
|
||||
}
|
||||
```
|
||||
|
||||
Design choices:
|
||||
|
||||
- **Agent index ownership**:
|
||||
- The server, not the model, owns id assignment.
|
||||
- The `agent_idx` input is treated as a hint and ignored; the actual id is returned in metadata:
|
||||
|
||||
```rust
|
||||
extra["agent_idx"] = json!(assigned_id);
|
||||
extra["parent_agent_idx"] = json!(caller_agent_id);
|
||||
```
|
||||
|
||||
- **Defaults**:
|
||||
- `context_strategy`: defaults to `New`.
|
||||
- `instructions`: defaults to the parent agent’s instructions (or session’s `developer_instructions` / `user_instructions`).
|
||||
- `sandbox_policy`, `model`: default to the session’s current config.
|
||||
- The child’s requested `sandbox_policy` is intended to be at least as strict as the parent/session; relaxing policies should result in an error.
|
||||
- TODO(jif): define and implement a formal “strictness” ordering for this policy and enforce it in `collaboration.init_agent`.
|
||||
|
||||
**Output**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationInitAgentMetadata {
|
||||
pub message_tool_call_success: bool,
|
||||
pub message_tool_call_error_should_penalize_model: bool,
|
||||
pub extra: ExtraMetadata,
|
||||
}
|
||||
|
||||
pub struct CollaborationInitAgentOutput {
|
||||
/// Status text from rollout_state.init_agent(...), e.g. success or error reason.
|
||||
pub content: String,
|
||||
pub metadata: CollaborationInitAgentMetadata,
|
||||
}
|
||||
```
|
||||
|
||||
Semantics:
|
||||
|
||||
- On success:
|
||||
- New `AgentState` is added to `CollaborationState.agents`.
|
||||
- `extra` includes at least:
|
||||
- `agent_idx`: assigned id.
|
||||
- `parent_agent_idx`: id of the caller.
|
||||
- `depth`: depth of the new agent.
|
||||
- `max_juice`: the agent’s juice cap.
|
||||
- `message_tool_call_success = true`.
|
||||
- On failure:
|
||||
- `content` is a concise error message (e.g., “max agent count reached”).
|
||||
- `message_tool_call_success = false`.
|
||||
- `message_tool_call_error_should_penalize_model`:
|
||||
- `true` for model mistakes (invalid indices, negative juice, unsupported combination).
|
||||
- `false` for environment limits (max agents / depth reached).
|
||||
|
||||
### 5.2 `collaboration.send`
|
||||
|
||||
**Purpose**
|
||||
|
||||
Send a textual message from the calling agent to one or more recipient agents.
|
||||
|
||||
**Input**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationSendInput {
|
||||
/// Indices of the agents to send the message to.
|
||||
/// Must be in [0, num_agents - 1] at runtime.
|
||||
pub recipients: Vec<u32>,
|
||||
/// The message to send.
|
||||
pub message: String,
|
||||
}
|
||||
```
|
||||
|
||||
Sender identification:
|
||||
|
||||
- The sender is the agent currently executing the tool call.
|
||||
- We track this via `extra` on tool calls:
|
||||
- When a tool call originates from a particular agent, we tag it with `extra["caller_agent_idx"]`.
|
||||
- Within the handler, we resolve that id from the turn’s collaboration context (see §6).
|
||||
|
||||
**Behavior**
|
||||
|
||||
For each `recipient`:
|
||||
|
||||
- Validate `recipient` against `CollaborationState.agents.len()`.
|
||||
- Append a `ResponseItem::Message { role: "user", content: … }` to the recipient’s `AgentState.history`, with content enriched to include the sender:
|
||||
|
||||
- E.g. `"From agent 0: <message>"`.
|
||||
|
||||
- Optionally, add a summarized `ResponseItem` into the session history so the UI can show “agent‑to‑agent” chatter.
|
||||
|
||||
**Output**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationSendMetadata {
|
||||
pub message_tool_call_success: bool,
|
||||
pub message_tool_call_error_should_penalize_model: bool,
|
||||
pub is_send_success_msg: Option<bool>,
|
||||
pub message_content_str: Option<String>,
|
||||
pub extra: ExtraMetadata,
|
||||
}
|
||||
|
||||
pub struct CollaborationSendOutput {
|
||||
pub content: String, // e.g. "Message sent successfully."
|
||||
pub metadata: CollaborationSendMetadata,
|
||||
}
|
||||
```
|
||||
|
||||
Conventions:
|
||||
|
||||
- On full success:
|
||||
- `is_send_success_msg = Some(true)`.
|
||||
- `message_content_str = Some(message.clone())`.
|
||||
- `extra["recipients"] = json!(recipients)`.
|
||||
- On partial success (some valid, some invalid ids):
|
||||
- `message_tool_call_success = false`.
|
||||
- `message_tool_call_error_should_penalize_model = true`.
|
||||
- Content explains which ids failed.
|
||||
|
||||
### 5.3 `collaboration.wait`
|
||||
|
||||
**Purpose**
|
||||
|
||||
Allow the main agent to yield control and synchronize with child agents that are already running in their own loops. This is where we *observe and bound* other agents’ progress, not where we start their streams.
|
||||
|
||||
**Input**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationWaitInput {
|
||||
/// Maximum duration to wait, measured in tokens. Must be >= 0.
|
||||
pub max_duration: u32,
|
||||
/// Optional list of child agents to wait for.
|
||||
/// Must be direct children of the caller.
|
||||
pub recipients: Option<Vec<u32>>,
|
||||
}
|
||||
```
|
||||
|
||||
Interpretation:
|
||||
|
||||
- `max_duration` is a soft token budget *for all agents during this wait call*.
|
||||
- Codex approximates this in terms of:
|
||||
- Incremental `TokenUsageInfo` per agent turn, using their per‑agent `ContextManager`.
|
||||
- And/or a fixed per‑turn granularity (e.g., “at most N turns across agents”).
|
||||
|
||||
Execution model:
|
||||
|
||||
- Collaboration runs as a dedicated `SessionTask` (a “supervisor”) which in turn spawns **one Tokio task per agent**:
|
||||
- Each agent task owns its own `run_task` loop, built from a per‑agent `TurnContext` and `ContextManager`.
|
||||
- These agent tasks continuously consume their model streams and write results into per‑agent histories and internal queues, so upstream servers never block on the parent.
|
||||
- `collaboration.wait` is implemented inside the supervisor task as:
|
||||
- A blocking point for the **main agent**: it observes per‑agent queues / histories and the total token deltas since the last wait.
|
||||
- If `recipients` is `Some`, it only considers those direct children; otherwise it considers *all* direct children of the caller.
|
||||
- It blocks until one of the following holds:
|
||||
- Enough new tokens have been consumed across the targeted agents (reaching `max_duration`).
|
||||
- All targeted agents are idle/completed, out of juice, or closed.
|
||||
- Cancellation is requested.
|
||||
- When `wait` returns, the main agent can inspect child results via `collaboration.get_state` without having had to manually drive each child’s `run_turn`.
|
||||
|
||||
**Output**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationWaitMetadata {
|
||||
pub message_tool_call_success: bool,
|
||||
pub message_tool_call_error_should_penalize_model: bool,
|
||||
pub is_wait_success_msg: Option<bool>,
|
||||
pub extra: ExtraMetadata,
|
||||
}
|
||||
|
||||
pub struct CollaborationWaitOutput {
|
||||
/// Text like "Started waiting successfully." or "Finished waiting."
|
||||
pub content: String,
|
||||
pub metadata: CollaborationWaitMetadata,
|
||||
}
|
||||
```
|
||||
|
||||
`extra` can include:
|
||||
|
||||
- A snapshot of which agents ran and for how many tokens:
|
||||
|
||||
- `extra["agents_ran"] = json!([{ "agent_idx": 1, "delta_tokens": 128, "status": "running" }, …])`.
|
||||
|
||||
- Whether token/juice limits were hit:
|
||||
|
||||
- `extra["token_budget_exhausted"] = json!(true)` when `max_duration` was the stopping reason.
|
||||
|
||||
### 5.4 `collaboration.get_state`
|
||||
|
||||
**Purpose**
|
||||
|
||||
Give the main agent a high‑level view of the collaboration graph (agents, statuses, juice, depth).
|
||||
|
||||
**Input/Output**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationGetStateInput {
|
||||
// No fields; the call takes an empty object.
|
||||
}
|
||||
|
||||
pub struct CollaborationGetStateMetadata {
|
||||
/// Always true on successful get_state.
|
||||
pub message_tool_call_success: Option<bool>,
|
||||
pub message_tool_call_error_should_penalize_model: Option<bool>,
|
||||
pub extra: ExtraMetadata,
|
||||
}
|
||||
|
||||
pub struct CollaborationGetStateOutput {
|
||||
/// Human-readable description of all agents' states.
|
||||
pub content: String,
|
||||
pub metadata: CollaborationGetStateMetadata,
|
||||
}
|
||||
```
|
||||
|
||||
Implementation:
|
||||
|
||||
- Walk `CollaborationState.agents` and construct:
|
||||
- A human‑readable table/list in `content` (suitable for the model to parse informally).
|
||||
- A structured mirror in `extra["agents"]`, e.g.:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"agent_idx": 0,
|
||||
"parent_agent_idx": null,
|
||||
"depth": 0,
|
||||
"status": "running",
|
||||
"used_juice": 512,
|
||||
"max_juice": 2000
|
||||
},
|
||||
...
|
||||
]
|
||||
```
|
||||
|
||||
- `message_tool_call_success = Some(true)` on success; errors are unlikely here.
|
||||
|
||||
`get_state` is the **primary** way the main agent inspects child progress and results. It is expected to grow richer over time (e.g., enums for lifecycle state, last result metadata) even if the top‑level schema remains stable.
|
||||
|
||||
### 5.5 `collaboration.close`
|
||||
|
||||
**Purpose**
|
||||
|
||||
Allow a parent agent to explicitly close one or more of its child agents. Closing an agent:
|
||||
|
||||
- Stops its Tokio task and prevents further model calls.
|
||||
- Recursively closes all of its descendants in the agent tree.
|
||||
- Marks their lifecycle state as `Closed` while keeping enough metadata for inspection via `get_state`.
|
||||
|
||||
**Input**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationCloseInput {
|
||||
/// Indices of the agents to close.
|
||||
/// Each must be a direct child of the caller.
|
||||
pub recipients: Vec<u32>,
|
||||
}
|
||||
```
|
||||
|
||||
**Output**
|
||||
|
||||
```rust
|
||||
pub struct CollaborationCloseMetadata {
|
||||
pub message_tool_call_success: bool,
|
||||
pub message_tool_call_error_should_penalize_model: bool,
|
||||
pub extra: ExtraMetadata,
|
||||
}
|
||||
|
||||
pub struct CollaborationCloseOutput {
|
||||
/// Human-readable status, e.g. "Closed 2 agents (and their descendants)."
|
||||
pub content: String,
|
||||
pub metadata: CollaborationCloseMetadata,
|
||||
}
|
||||
```
|
||||
|
||||
Semantics:
|
||||
|
||||
- On success:
|
||||
- For each `recipient` that is a direct child of the caller:
|
||||
- Cancel its Tokio task and mark it `Closed`.
|
||||
- Recursively do the same for all its descendants using the `children` adjacency.
|
||||
- `extra` may include:
|
||||
- `closed_agent_indices`: all agents that were closed.
|
||||
- `message_tool_call_success = true`.
|
||||
- On failure:
|
||||
- If any `recipient` is not a direct child of the caller (or refers to the root), reject the call:
|
||||
- `message_tool_call_success = false`.
|
||||
- `message_tool_call_error_should_penalize_model = true`.
|
||||
- `content` explains the violation.
|
||||
|
||||
---
|
||||
|
||||
## 6. Mapping Tool Calls to Agents
|
||||
|
||||
Because all agents share the same tool ecosystem, we need a way to know *which agent* is currently issuing a tool call.
|
||||
|
||||
Approach:
|
||||
|
||||
- Associate each `TurnContext` / `sub_id` with an `AgentId` inside `CollaborationState`:
|
||||
- When the supervisor spawns an agent’s Tokio task, it registers a mapping from that agent’s `sub_id` (or a dedicated identifier in the `TurnContext`) to `AgentId`.
|
||||
- All model streams and tool calls originating from that agent run under that `TurnContext`.
|
||||
- When building tool calls in `ToolRouter::build_tool_call`, we extend `ToolPayload` or `ToolInvocation` with:
|
||||
|
||||
- `caller_agent_idx: AgentId` (looked up from the current `TurnContext`’s identifier).
|
||||
|
||||
- The collaboration tool handlers read `caller_agent_idx` to:
|
||||
- Determine parent/child relationships for `init_agent`.
|
||||
- Record messages in the correct `AgentState.history` for `send`.
|
||||
- Attribute juice consumption to the right agent for `wait` and enforce that parents only target direct children.
|
||||
|
||||
This separation keeps the rest of the tool ecosystem unchanged: non‑collaboration tools are unaware of agents and continue to operate exactly as before, but their events are recorded in the appropriate agent’s history when run on behalf of that agent.
|
||||
|
||||
---
|
||||
|
||||
## 7. Limits, Safety, and Error Handling
|
||||
|
||||
### 7.1 Max Agents and Depth
|
||||
|
||||
Enforced in `CollaborationState`:
|
||||
|
||||
- On `init_agent`:
|
||||
- If `agents.len() >= max_agents` → reject with a clear error string and `message_tool_call_success = false`.
|
||||
- If `parent.depth + 1 > max_depth` → reject similarly.
|
||||
- The main agent (id 0) always exists once collaboration is used and does not count towards depth.
|
||||
|
||||
### 7.2 Juice and Token Accounting
|
||||
|
||||
- Per agent:
|
||||
- `used_juice` is derived from that agent’s `ContextManager::get_total_token_usage()`.
|
||||
- `max_juice` comes from the `juice` field on `CollaborationInitAgentInput`.
|
||||
- After each agent turn, we compute the delta and update `used_juice`.
|
||||
- Once `used_juice >= max_juice`, the agent’s status is set to `AgentLifecycleState::Exhausted` and `wait` will no longer schedule it.
|
||||
- Per `wait` call:
|
||||
- Track a separate counter for the duration of that call based on token deltas across all agents that ran.
|
||||
|
||||
### 7.3 Error Surfaces
|
||||
|
||||
Examples:
|
||||
|
||||
- Invalid agent indices:
|
||||
- `init_agent`: if the parent agent id is invalid, this is a model bug; penalize.
|
||||
- `send`: out‑of‑bounds recipients; penalize.
|
||||
- `wait`/`get_state`: essentially infallible; errors here are environment errors.
|
||||
- Limit violations:
|
||||
- Max agents / depth / juice: treat as environment constraints; do **not** penalize the model but describe limits clearly in the error content.
|
||||
|
||||
---
|
||||
|
||||
## 8. Integration with Existing Components
|
||||
|
||||
### 8.1 `Session` and `SessionState`
|
||||
|
||||
Add collaboration‑related helpers on `Session`:
|
||||
|
||||
- `fn collaboration(&self) -> &CollaborationStateHandle` to access the shared state (behind a mutex).
|
||||
- Helper methods to:
|
||||
- Initialize the root agent from the current session config.
|
||||
- Get the `AgentState` for the “main” agent (id 0).
|
||||
- Map `sub_id` or tool call metadata to `AgentId`.
|
||||
|
||||
`SessionState` can remain focused on global history and token/rate limit tracking; collaboration uses `ContextManager` in parallel for per‑agent history.
|
||||
|
||||
### 8.2 `run_task` / `run_turn`
|
||||
|
||||
The single‑agent path (today’s behavior) remains:
|
||||
|
||||
- `RegularTask::run` → `run_task` loop.
|
||||
- `run_task` builds `prompt` from *session* history and calls `run_turn`.
|
||||
|
||||
For collaboration, we introduce a **CollaborationTask**:
|
||||
|
||||
- A specialized `SessionTask` that:
|
||||
- Runs the main agent’s `run_task` loop (the one that can call `collaboration.*` tools).
|
||||
- Spawns and supervises one Tokio task per child agent, each running its own `run_task` / `run_turn` loop with a per‑agent `TurnContext` and `ContextManager`.
|
||||
- Maintains per‑agent queues and histories so the main agent can retrieve child outputs cheaply.
|
||||
- `collaboration.wait` is implemented within this CollaborationTask as a synchronization point with those per‑agent tasks; the agents themselves do not depend on `wait` to consume their model streams.
|
||||
|
||||
We still **reuse** the existing `run_turn` and streaming/tool machinery: collaboration only changes how many `run_task` loops we run in parallel and how we coordinate them, not the semantics of a single agent’s interaction with the model.
|
||||
|
||||
### 8.3 Tooling and Registry
|
||||
|
||||
We add collaboration handlers under `core/src/tools/handlers/collaboration.rs`:
|
||||
|
||||
- Define `ToolSpec`s for the four tools, matching the Rust structs defined above (via `JsonSchema`).
|
||||
- Implement `ToolHandler` for each, reading/writing `CollaborationState`, `Session`, and `TurnContext`, and returning structured outputs.
|
||||
|
||||
The `ToolRouter` remains unchanged; collaboration is just another tool namespace (`"collaboration.init_agent"`, etc.).
|
||||
|
||||
---
|
||||
|
||||
## 9. Rollout and Telemetry
|
||||
|
||||
Because all agent interactions still happen within a single `Session`:
|
||||
|
||||
- **Rollout recorder**:
|
||||
- We continue to record `RolloutItem::ResponseItem` for:
|
||||
- Messages from the main agent.
|
||||
- Summaries of child agents’ activity (collaboration events).
|
||||
- Per‑agent histories live alongside this and can be serialized separately if needed for debugging.
|
||||
- **Telemetry (OTel)**:
|
||||
- For child agents, we re‑use the existing `OtelEventManager`, potentially tagging:
|
||||
- `agent_idx`, `parent_agent_idx`, `depth` as attributes on tool and model spans.
|
||||
|
||||
This gives us visibility into:
|
||||
|
||||
- How many agents are spawned per session.
|
||||
- How much juice they consume.
|
||||
- How collaboration changes overall token usage and user outcomes.
|
||||
|
||||
---
|
||||
|
||||
## 10. Implementation Outline
|
||||
|
||||
1. **Core types and state**
|
||||
- Introduce `AgentId`, `AgentState`, `CollaborationLimits`, `CollaborationState`.
|
||||
- Add a collaboration handle to `Session` (with appropriate locking).
|
||||
- Implement methods for creating the root agent and looking up agents.
|
||||
2. **Per‑agent configuration**
|
||||
- Extend `SessionConfiguration` cloning logic to produce per‑agent configs.
|
||||
- Implement `make_agent_turn_context` helpers.
|
||||
3. **Tool specs and handlers**
|
||||
- Add `collaboration.rs` handler module.
|
||||
- Define `ToolSpec`s for `collaboration.init_agent`, `.send`, `.wait`, `.get_state`, `.close`.
|
||||
- Implement tool handlers using the structs in the prompt and `ExtraMetadata`.
|
||||
4. **Agent scheduling**
|
||||
- Implement a `CollaborationTask` supervisor that spawns one Tokio task per agent, tracks per‑agent token deltas/juice, and implements `collaboration.wait` as a synchronization barrier over those tasks rather than directly calling `run_turn`.
|
||||
5. **History integration**
|
||||
- Decide on a minimal, stable projection of agent activity into `SessionState.history` (e.g., short assistant messages summarizing agent actions).
|
||||
- Ensure per‑agent histories remain valid `ContextManager`s (respecting call/output invariants).
|
||||
6. **Limits and validation**
|
||||
- Enforce `max_agents`, `max_depth`, and per‑agent `max_juice`.
|
||||
- Wire `message_tool_call_error_should_penalize_model` for model vs environment errors.
|
||||
|
||||
This design keeps the existing Codex single‑agent loop intact while introducing a structured, tool‑driven multi‑agent collaboration model with clear state boundaries and safety limits.
|
||||
+264
-46
@@ -105,8 +105,13 @@ use crate::rollout::RolloutRecorderParams;
|
||||
use crate::rollout::map_session_init_error;
|
||||
use crate::shell;
|
||||
use crate::state::ActiveTurn;
|
||||
use crate::state::AgentId;
|
||||
use crate::state::AgentState;
|
||||
use crate::state::CollaborationLimits;
|
||||
use crate::state::CollaborationState;
|
||||
use crate::state::SessionServices;
|
||||
use crate::state::SessionState;
|
||||
use crate::tasks::CollaborationSupervisor;
|
||||
use crate::tasks::GhostSnapshotTask;
|
||||
use crate::tasks::ReviewTask;
|
||||
use crate::tasks::SessionTask;
|
||||
@@ -265,6 +270,8 @@ pub(crate) struct Session {
|
||||
state: Mutex<SessionState>,
|
||||
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
|
||||
pub(crate) services: SessionServices,
|
||||
collaboration: Mutex<CollaborationState>,
|
||||
collaboration_supervisor: Mutex<Option<CollaborationSupervisor>>,
|
||||
next_internal_sub_id: AtomicU64,
|
||||
}
|
||||
|
||||
@@ -290,6 +297,7 @@ pub(crate) struct TurnContext {
|
||||
pub(crate) tool_call_gate: Arc<ReadinessFlag>,
|
||||
pub(crate) exec_policy: Arc<ExecPolicy>,
|
||||
pub(crate) truncation_policy: TruncationPolicy,
|
||||
pub(crate) collaboration_agent: AgentId,
|
||||
}
|
||||
|
||||
impl TurnContext {
|
||||
@@ -304,9 +312,13 @@ impl TurnContext {
|
||||
.as_deref()
|
||||
.unwrap_or(compact::SUMMARIZATION_PROMPT)
|
||||
}
|
||||
|
||||
pub(crate) fn collaboration_agent(&self) -> AgentId {
|
||||
self.collaboration_agent
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct SessionConfiguration {
|
||||
/// Provider identifier ("openai", "openrouter", ...).
|
||||
provider: ModelProviderInfo,
|
||||
@@ -377,6 +389,22 @@ impl SessionConfiguration {
|
||||
}
|
||||
next_configuration
|
||||
}
|
||||
|
||||
pub(crate) fn developer_instructions(&self) -> Option<String> {
|
||||
self.developer_instructions.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn user_instructions(&self) -> Option<String> {
|
||||
self.user_instructions.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn approval_policy(&self) -> AskForApproval {
|
||||
self.approval_policy
|
||||
}
|
||||
|
||||
pub(crate) fn sandbox_policy(&self) -> SandboxPolicy {
|
||||
self.sandbox_policy.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default, Clone)]
|
||||
@@ -398,6 +426,7 @@ impl Session {
|
||||
session_configuration: &SessionConfiguration,
|
||||
conversation_id: ConversationId,
|
||||
sub_id: String,
|
||||
agent_id: AgentId,
|
||||
) -> TurnContext {
|
||||
let config = session_configuration.original_config_do_not_use.clone();
|
||||
let model_family = find_family_for_model(&session_configuration.model)
|
||||
@@ -449,9 +478,26 @@ impl Session {
|
||||
tool_call_gate: Arc::new(ReadinessFlag::new()),
|
||||
exec_policy: session_configuration.exec_policy.clone(),
|
||||
truncation_policy: TruncationPolicy::new(&per_turn_config),
|
||||
collaboration_agent: agent_id,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn make_collaboration_turn_context(
|
||||
&self,
|
||||
agent: &AgentState,
|
||||
sub_id: String,
|
||||
) -> TurnContext {
|
||||
Self::make_turn_context(
|
||||
Some(Arc::clone(&self.services.auth_manager)),
|
||||
&self.services.otel_event_manager,
|
||||
agent.config.provider.clone(),
|
||||
&agent.config,
|
||||
self.conversation_id,
|
||||
sub_id,
|
||||
agent.id,
|
||||
)
|
||||
}
|
||||
|
||||
async fn new(
|
||||
session_configuration: SessionConfiguration,
|
||||
config: Arc<Config>,
|
||||
@@ -556,6 +602,7 @@ impl Session {
|
||||
|
||||
// Create the mutable state for the Session.
|
||||
let state = SessionState::new(session_configuration.clone());
|
||||
let collaboration = CollaborationState::new(CollaborationLimits::default());
|
||||
|
||||
let services = SessionServices {
|
||||
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
|
||||
@@ -576,6 +623,8 @@ impl Session {
|
||||
state: Mutex::new(state),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
collaboration: Mutex::new(collaboration),
|
||||
collaboration_supervisor: Mutex::new(None),
|
||||
next_internal_sub_id: AtomicU64::new(0),
|
||||
});
|
||||
|
||||
@@ -642,6 +691,22 @@ impl Session {
|
||||
self.tx_event.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn collaboration_state(&self) -> &Mutex<CollaborationState> {
|
||||
&self.collaboration
|
||||
}
|
||||
|
||||
pub(crate) async fn ensure_collaboration_supervisor(
|
||||
self: &Arc<Self>,
|
||||
) -> CollaborationSupervisor {
|
||||
let mut guard = self.collaboration_supervisor.lock().await;
|
||||
if let Some(existing) = guard.clone() {
|
||||
return existing;
|
||||
}
|
||||
let handle = CollaborationSupervisor::spawn(Arc::clone(self));
|
||||
*guard = Some(handle.clone());
|
||||
handle
|
||||
}
|
||||
|
||||
/// Ensure all rollout writes are durably flushed.
|
||||
pub(crate) async fn flush_rollout(&self) {
|
||||
let recorder = {
|
||||
@@ -757,13 +822,31 @@ impl Session {
|
||||
&session_configuration,
|
||||
self.conversation_id,
|
||||
sub_id,
|
||||
AgentId(0),
|
||||
);
|
||||
if let Some(final_schema) = updates.final_output_json_schema {
|
||||
turn_context.final_output_json_schema = final_schema;
|
||||
}
|
||||
|
||||
let session_history = self.clone_history().await;
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
collab.ensure_root_agent(&session_configuration, &session_history);
|
||||
drop(collab);
|
||||
self.register_sub_id(AgentId(0), turn_context.sub_id.clone())
|
||||
.await;
|
||||
Arc::new(turn_context)
|
||||
}
|
||||
|
||||
pub(crate) async fn register_sub_id(&self, agent: AgentId, sub_id: String) {
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
collab.register_sub_id(agent, sub_id);
|
||||
}
|
||||
|
||||
pub(crate) async fn current_session_configuration(&self) -> SessionConfiguration {
|
||||
let state = self.state.lock().await;
|
||||
state.session_configuration.clone()
|
||||
}
|
||||
|
||||
fn build_environment_update_item(
|
||||
&self,
|
||||
previous: Option<&Arc<TurnContext>>,
|
||||
@@ -1032,8 +1115,32 @@ impl Session {
|
||||
items: &[ResponseItem],
|
||||
turn_context: &TurnContext,
|
||||
) {
|
||||
let mut state = self.state.lock().await;
|
||||
state.record_items(items.iter(), turn_context.truncation_policy);
|
||||
let agent = turn_context.collaboration_agent();
|
||||
if agent == AgentId(0) {
|
||||
let (history, token_info, config) = {
|
||||
let mut state = self.state.lock().await;
|
||||
state.record_items(items.iter(), turn_context.truncation_policy);
|
||||
(
|
||||
state.clone_history(),
|
||||
state.token_info(),
|
||||
state.session_configuration.clone(),
|
||||
)
|
||||
};
|
||||
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
collab.ensure_root_agent(&config, &history);
|
||||
if let Some(root) = collab.agent_mut(AgentId(0)) {
|
||||
root.history = history;
|
||||
root.history.set_token_info(token_info);
|
||||
}
|
||||
} else {
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(agent_state) = collab.agent_mut(agent) {
|
||||
agent_state
|
||||
.history
|
||||
.record_items(items.iter(), turn_context.truncation_policy);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn record_model_warning(&self, message: impl Into<String>, ctx: &TurnContext) {
|
||||
@@ -1053,8 +1160,7 @@ impl Session {
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_history(&self, items: Vec<ResponseItem>) {
|
||||
let mut state = self.state.lock().await;
|
||||
state.replace_history(items);
|
||||
self.set_history_for_agent(AgentId(0), items, None).await;
|
||||
}
|
||||
|
||||
async fn persist_rollout_response_items(&self, items: &[ResponseItem]) {
|
||||
@@ -1121,8 +1227,34 @@ impl Session {
|
||||
}
|
||||
|
||||
pub(crate) async fn clone_history(&self) -> ContextManager {
|
||||
let state = self.state.lock().await;
|
||||
state.clone_history()
|
||||
self.clone_history_for_agent(AgentId(0)).await
|
||||
}
|
||||
|
||||
pub(crate) async fn clone_history_for_agent(&self, agent: AgentId) -> ContextManager {
|
||||
if agent == AgentId(0) {
|
||||
let state = self.state.lock().await;
|
||||
return state.clone_history();
|
||||
}
|
||||
let collab = self.collaboration.lock().await;
|
||||
collab
|
||||
.clone_agent_history(agent)
|
||||
.unwrap_or_else(ContextManager::new)
|
||||
}
|
||||
|
||||
pub(crate) async fn set_history_for_agent(
|
||||
&self,
|
||||
agent: AgentId,
|
||||
items: Vec<ResponseItem>,
|
||||
token_info: Option<TokenUsageInfo>,
|
||||
) {
|
||||
if agent == AgentId(0) {
|
||||
let mut state = self.state.lock().await;
|
||||
state.replace_history(items.clone());
|
||||
state.set_token_info(token_info.clone());
|
||||
}
|
||||
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
let _ = collab.set_agent_history(agent, items, token_info);
|
||||
}
|
||||
|
||||
pub(crate) async fn update_token_usage_info(
|
||||
@@ -1130,49 +1262,94 @@ impl Session {
|
||||
turn_context: &TurnContext,
|
||||
token_usage: Option<&TokenUsage>,
|
||||
) {
|
||||
{
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(token_usage) = token_usage {
|
||||
state.update_token_info_from_usage(
|
||||
token_usage,
|
||||
turn_context.client.get_model_context_window(),
|
||||
);
|
||||
let agent = turn_context.collaboration_agent();
|
||||
if agent == AgentId(0) {
|
||||
let token_info = {
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(token_usage) = token_usage {
|
||||
state.update_token_info_from_usage(
|
||||
token_usage,
|
||||
turn_context.client.get_model_context_window(),
|
||||
);
|
||||
}
|
||||
state.token_info()
|
||||
};
|
||||
{
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(root) = collab.agent_mut(AgentId(0)) {
|
||||
root.history.set_token_info(token_info);
|
||||
}
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
} else if let Some(token_usage) = token_usage {
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(agent_state) = collab.agent_mut(agent) {
|
||||
agent_state
|
||||
.history
|
||||
.update_token_info(token_usage, turn_context.client.get_model_context_window());
|
||||
}
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) {
|
||||
let Some(estimated_total_tokens) = self
|
||||
.clone_history()
|
||||
.await
|
||||
.estimate_token_count(turn_context)
|
||||
else {
|
||||
let agent = turn_context.collaboration_agent();
|
||||
let history = self.clone_history_for_agent(agent).await;
|
||||
let Some(estimated_total_tokens) = history.estimate_token_count(turn_context) else {
|
||||
return;
|
||||
};
|
||||
{
|
||||
let mut state = self.state.lock().await;
|
||||
let mut info = state.token_info().unwrap_or(TokenUsageInfo {
|
||||
total_token_usage: TokenUsage::default(),
|
||||
last_token_usage: TokenUsage::default(),
|
||||
model_context_window: None,
|
||||
});
|
||||
|
||||
info.last_token_usage = TokenUsage {
|
||||
input_tokens: 0,
|
||||
cached_input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_output_tokens: 0,
|
||||
total_tokens: estimated_total_tokens.max(0),
|
||||
if agent == AgentId(0) {
|
||||
let token_info = {
|
||||
let mut state = self.state.lock().await;
|
||||
let mut info = state.token_info().unwrap_or(TokenUsageInfo {
|
||||
total_token_usage: TokenUsage::default(),
|
||||
last_token_usage: TokenUsage::default(),
|
||||
model_context_window: None,
|
||||
});
|
||||
|
||||
info.last_token_usage = TokenUsage {
|
||||
input_tokens: 0,
|
||||
cached_input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_output_tokens: 0,
|
||||
total_tokens: estimated_total_tokens.max(0),
|
||||
};
|
||||
|
||||
if info.model_context_window.is_none() {
|
||||
info.model_context_window = turn_context.client.get_model_context_window();
|
||||
}
|
||||
|
||||
state.set_token_info(Some(info.clone()));
|
||||
info
|
||||
};
|
||||
|
||||
if info.model_context_window.is_none() {
|
||||
info.model_context_window = turn_context.client.get_model_context_window();
|
||||
{
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(root) = collab.agent_mut(AgentId(0)) {
|
||||
root.history.set_token_info(Some(token_info));
|
||||
}
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
} else {
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(agent_state) = collab.agent_mut(agent) {
|
||||
let mut info = agent_state.history.token_info().unwrap_or(TokenUsageInfo {
|
||||
total_token_usage: TokenUsage::default(),
|
||||
last_token_usage: TokenUsage::default(),
|
||||
model_context_window: None,
|
||||
});
|
||||
info.last_token_usage = TokenUsage {
|
||||
input_tokens: 0,
|
||||
cached_input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
reasoning_output_tokens: 0,
|
||||
total_tokens: estimated_total_tokens.max(0),
|
||||
};
|
||||
if info.model_context_window.is_none() {
|
||||
info.model_context_window = turn_context.client.get_model_context_window();
|
||||
}
|
||||
agent_state.history.set_token_info(Some(info));
|
||||
}
|
||||
|
||||
state.set_token_info(Some(info));
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn update_rate_limits(
|
||||
@@ -1197,13 +1374,28 @@ impl Session {
|
||||
}
|
||||
|
||||
pub(crate) async fn set_total_tokens_full(&self, turn_context: &TurnContext) {
|
||||
let agent = turn_context.collaboration_agent();
|
||||
let context_window = turn_context.client.get_model_context_window();
|
||||
if let Some(context_window) = context_window {
|
||||
{
|
||||
let mut state = self.state.lock().await;
|
||||
state.set_token_usage_full(context_window);
|
||||
if agent == AgentId(0) {
|
||||
let token_info = {
|
||||
let mut state = self.state.lock().await;
|
||||
state.set_token_usage_full(context_window);
|
||||
state.token_info()
|
||||
};
|
||||
{
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(root) = collab.agent_mut(AgentId(0)) {
|
||||
root.history.set_token_info(token_info);
|
||||
}
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
} else {
|
||||
let mut collab = self.collaboration.lock().await;
|
||||
if let Some(agent_state) = collab.agent_mut(agent) {
|
||||
agent_state.history.set_token_usage_full(context_window);
|
||||
}
|
||||
}
|
||||
self.send_token_count_event(turn_context).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1494,8 +1686,8 @@ mod handlers {
|
||||
use crate::config::Config;
|
||||
use crate::mcp::auth::compute_auth_statuses;
|
||||
use crate::review_prompts::resolve_review_request;
|
||||
use crate::tasks::CollaborationTask;
|
||||
use crate::tasks::CompactTask;
|
||||
use crate::tasks::RegularTask;
|
||||
use crate::tasks::UndoTask;
|
||||
use crate::tasks::UserShellCommandTask;
|
||||
use codex_protocol::custom_prompts::CustomPrompt;
|
||||
@@ -1572,7 +1764,7 @@ mod handlers {
|
||||
.await;
|
||||
}
|
||||
|
||||
sess.spawn_task(Arc::clone(¤t_context), items, RegularTask)
|
||||
sess.spawn_task(Arc::clone(¤t_context), items, CollaborationTask)
|
||||
.await;
|
||||
*previous_context = Some(current_context);
|
||||
}
|
||||
@@ -1901,6 +2093,7 @@ async fn spawn_review_thread(
|
||||
tool_call_gate: Arc::new(ReadinessFlag::new()),
|
||||
exec_policy: parent_turn_context.exec_policy.clone(),
|
||||
truncation_policy: TruncationPolicy::new(&per_turn_config),
|
||||
collaboration_agent: AgentId(0),
|
||||
};
|
||||
|
||||
// Seed the child task with the review prompt as the initial user message.
|
||||
@@ -2179,6 +2372,23 @@ async fn run_turn(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn run_collaboration_turn(
|
||||
sess: Arc<Session>,
|
||||
turn_context: Arc<TurnContext>,
|
||||
turn_diff_tracker: SharedTurnDiffTracker,
|
||||
input: Vec<ResponseItem>,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> CodexResult<Vec<ProcessedResponseItem>> {
|
||||
run_turn(
|
||||
sess,
|
||||
turn_context,
|
||||
turn_diff_tracker,
|
||||
input,
|
||||
cancellation_token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// When the model is prompted, it returns a stream of events. Some of these
|
||||
/// events map to a `ResponseItem`. A `ResponseItem` may need to be
|
||||
/// "handled" such that it produces a `ResponseInputItem` that needs to be
|
||||
@@ -2444,7 +2654,7 @@ async fn handle_non_tool_response_item(item: &ResponseItem) -> Option<TurnItem>
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<String> {
|
||||
pub(crate) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<String> {
|
||||
responses.iter().rev().find_map(|item| {
|
||||
if let ResponseItem::Message { role, content, .. } = item {
|
||||
if role == "assistant" {
|
||||
@@ -2782,6 +2992,7 @@ mod tests {
|
||||
};
|
||||
|
||||
let state = SessionState::new(session_configuration.clone());
|
||||
let collaboration = CollaborationState::new(CollaborationLimits::default());
|
||||
|
||||
let services = SessionServices {
|
||||
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
|
||||
@@ -2803,6 +3014,7 @@ mod tests {
|
||||
&session_configuration,
|
||||
conversation_id,
|
||||
"turn_id".to_string(),
|
||||
AgentId(0),
|
||||
);
|
||||
|
||||
let session = Session {
|
||||
@@ -2811,6 +3023,8 @@ mod tests {
|
||||
state: Mutex::new(state),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
collaboration: Mutex::new(collaboration),
|
||||
collaboration_supervisor: Mutex::new(None),
|
||||
next_internal_sub_id: AtomicU64::new(0),
|
||||
};
|
||||
|
||||
@@ -2860,6 +3074,7 @@ mod tests {
|
||||
};
|
||||
|
||||
let state = SessionState::new(session_configuration.clone());
|
||||
let collaboration = CollaborationState::new(CollaborationLimits::default());
|
||||
|
||||
let services = SessionServices {
|
||||
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
|
||||
@@ -2881,6 +3096,7 @@ mod tests {
|
||||
&session_configuration,
|
||||
conversation_id,
|
||||
"turn_id".to_string(),
|
||||
AgentId(0),
|
||||
));
|
||||
|
||||
let session = Arc::new(Session {
|
||||
@@ -2889,6 +3105,8 @@ mod tests {
|
||||
state: Mutex::new(state),
|
||||
active_turn: Mutex::new(None),
|
||||
services,
|
||||
collaboration: Mutex::new(collaboration),
|
||||
collaboration_supervisor: Mutex::new(None),
|
||||
next_internal_sub_id: AtomicU64::new(0),
|
||||
});
|
||||
|
||||
|
||||
@@ -590,7 +590,7 @@ mod tests {
|
||||
assert_eq!(entries.len(), 1);
|
||||
assert_eq!(entries[0].text, long_entry);
|
||||
|
||||
let pruned_len = std::fs::metadata(&history_path).expect("metadata").len() as u64;
|
||||
let pruned_len = std::fs::metadata(&history_path).expect("metadata").len();
|
||||
let max_bytes = config
|
||||
.history
|
||||
.max_bytes
|
||||
|
||||
@@ -195,6 +195,13 @@ pub fn find_family_for_model(slug: &str) -> Option<ModelFamily> {
|
||||
reasoning_summary_format: ReasoningSummaryFormat::Experimental,
|
||||
base_instructions: GPT_5_1_CODEX_MAX_INSTRUCTIONS.to_string(),
|
||||
apply_patch_tool_type: Some(ApplyPatchToolType::Freeform),
|
||||
experimental_supported_tools: vec![
|
||||
"grep_files".to_string(),
|
||||
"list_dir".to_string(),
|
||||
"read_file".to_string(),
|
||||
"test_sync_tool".to_string(),
|
||||
"collaboration".to_string(),
|
||||
],
|
||||
shell_type: ConfigShellToolType::ShellCommand,
|
||||
supports_parallel_tool_calls: true,
|
||||
support_verbosity: false,
|
||||
@@ -210,6 +217,12 @@ pub fn find_family_for_model(slug: &str) -> Option<ModelFamily> {
|
||||
reasoning_summary_format: ReasoningSummaryFormat::Experimental,
|
||||
base_instructions: GPT_5_CODEX_INSTRUCTIONS.to_string(),
|
||||
apply_patch_tool_type: Some(ApplyPatchToolType::Freeform),
|
||||
experimental_supported_tools: vec![
|
||||
"grep_files".to_string(),
|
||||
"list_dir".to_string(),
|
||||
"read_file".to_string(),
|
||||
"test_sync_tool".to_string(),
|
||||
],
|
||||
shell_type: ConfigShellToolType::ShellCommand,
|
||||
supports_parallel_tool_calls: true,
|
||||
support_verbosity: false,
|
||||
|
||||
@@ -0,0 +1,278 @@
|
||||
//! Session-scoped collaboration state for multi-agent flows.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::codex::SessionConfiguration;
|
||||
use crate::context_manager::ContextManager;
|
||||
use crate::protocol::TokenUsageInfo;
|
||||
use crate::truncate::TruncationPolicy;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
|
||||
pub(crate) struct AgentId(pub i32);
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum AgentLifecycleState {
|
||||
Idle { last_agent_message: Option<String> },
|
||||
Running,
|
||||
Exhausted,
|
||||
Error { error: String },
|
||||
Closed,
|
||||
WaitingForApproval { request: String },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ContextStrategy {
|
||||
New,
|
||||
Fork,
|
||||
Replace { history: Vec<ResponseItem> },
|
||||
}
|
||||
|
||||
impl Default for ContextStrategy {
|
||||
fn default() -> Self {
|
||||
Self::New
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct AgentState {
|
||||
pub(crate) id: AgentId,
|
||||
pub(crate) parent: Option<AgentId>,
|
||||
pub(crate) depth: i32,
|
||||
pub(crate) config: SessionConfiguration,
|
||||
pub(crate) instructions: Option<String>,
|
||||
pub(crate) status: AgentLifecycleState,
|
||||
pub(crate) history: ContextManager,
|
||||
}
|
||||
|
||||
impl AgentState {
|
||||
pub(crate) fn new_root(
|
||||
config: SessionConfiguration,
|
||||
history: ContextManager,
|
||||
instructions: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id: AgentId(0),
|
||||
parent: None,
|
||||
depth: 0,
|
||||
config,
|
||||
instructions,
|
||||
status: AgentLifecycleState::Idle {
|
||||
last_agent_message: None,
|
||||
},
|
||||
history,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn new_child(
|
||||
id: AgentId,
|
||||
parent: AgentId,
|
||||
depth: i32,
|
||||
config: SessionConfiguration,
|
||||
instructions: Option<String>,
|
||||
history: ContextManager,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
parent: Some(parent),
|
||||
depth,
|
||||
config,
|
||||
instructions,
|
||||
status: AgentLifecycleState::Idle {
|
||||
last_agent_message: None,
|
||||
},
|
||||
history,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(crate) struct CollaborationLimits {
|
||||
pub(crate) max_agents: i32,
|
||||
pub(crate) max_depth: i32,
|
||||
}
|
||||
|
||||
impl Default for CollaborationLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_agents: 8,
|
||||
max_depth: 4,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct CollaborationState {
|
||||
agents: Vec<AgentState>,
|
||||
children: HashMap<AgentId, Vec<AgentId>>,
|
||||
limits: CollaborationLimits,
|
||||
next_sub_id: i64,
|
||||
sub_ids: HashMap<String, AgentId>,
|
||||
}
|
||||
|
||||
impl CollaborationState {
|
||||
pub(crate) fn new(limits: CollaborationLimits) -> Self {
|
||||
Self {
|
||||
agents: Vec::new(),
|
||||
children: HashMap::new(),
|
||||
limits,
|
||||
next_sub_id: 0,
|
||||
sub_ids: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn limits(&self) -> &CollaborationLimits {
|
||||
&self.limits
|
||||
}
|
||||
|
||||
pub(crate) fn ensure_root_agent(
|
||||
&mut self,
|
||||
session_configuration: &SessionConfiguration,
|
||||
session_history: &ContextManager,
|
||||
) -> AgentId {
|
||||
if self.agents.is_empty() {
|
||||
let root = AgentState::new_root(
|
||||
session_configuration.clone(),
|
||||
session_history.clone(),
|
||||
session_configuration
|
||||
.developer_instructions()
|
||||
.or_else(|| session_configuration.user_instructions()),
|
||||
);
|
||||
self.agents.push(root);
|
||||
} else if let Some(root) = self.agents.get_mut(0) {
|
||||
root.config = session_configuration.clone();
|
||||
root.history = session_history.clone();
|
||||
if root.instructions.is_none() {
|
||||
root.instructions = session_configuration
|
||||
.developer_instructions()
|
||||
.or_else(|| session_configuration.user_instructions());
|
||||
}
|
||||
}
|
||||
AgentId(0)
|
||||
}
|
||||
|
||||
pub(crate) fn agents(&self) -> &[AgentState] {
|
||||
&self.agents
|
||||
}
|
||||
|
||||
pub(crate) fn agent(&self, id: AgentId) -> Option<&AgentState> {
|
||||
self.index_for(id).and_then(|idx| self.agents.get(idx))
|
||||
}
|
||||
|
||||
pub(crate) fn agent_mut(&mut self, id: AgentId) -> Option<&mut AgentState> {
|
||||
let index = self.index_for(id)?;
|
||||
self.agents.get_mut(index)
|
||||
}
|
||||
|
||||
pub(crate) fn next_agent_id(&self) -> AgentId {
|
||||
AgentId(self.agents.len() as i32)
|
||||
}
|
||||
|
||||
pub(crate) fn clone_agent_history(&self, id: AgentId) -> Option<ContextManager> {
|
||||
self.agent(id).map(|agent| agent.history.clone())
|
||||
}
|
||||
|
||||
pub(crate) fn set_agent_history(
|
||||
&mut self,
|
||||
id: AgentId,
|
||||
items: Vec<ResponseItem>,
|
||||
token_info: Option<TokenUsageInfo>,
|
||||
) -> Result<(), String> {
|
||||
let agent = self
|
||||
.agent_mut(id)
|
||||
.ok_or_else(|| format!("unknown agent {}", id.0))?;
|
||||
agent.history.replace(items);
|
||||
agent.history.set_token_info(token_info);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn record_message_for_agent(&mut self, id: AgentId, message: ResponseItem) {
|
||||
if let Some(agent) = self.agent_mut(id) {
|
||||
agent
|
||||
.history
|
||||
.record_items([message].iter(), TruncationPolicy::Bytes(10_000));
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn record_items_for_agent(
|
||||
&mut self,
|
||||
id: AgentId,
|
||||
items: &[ResponseItem],
|
||||
policy: TruncationPolicy,
|
||||
) {
|
||||
if let Some(agent) = self.agent_mut(id) {
|
||||
agent.history.record_items(items.iter(), policy);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn add_child(&mut self, mut agent: AgentState) -> Result<AgentId, String> {
|
||||
if self.agents.len() as i32 >= self.limits.max_agents {
|
||||
return Err("max agent count reached".to_string());
|
||||
}
|
||||
if agent.depth > self.limits.max_depth {
|
||||
return Err("max collaboration depth reached".to_string());
|
||||
}
|
||||
|
||||
let id = self.next_agent_id();
|
||||
agent.id = id;
|
||||
|
||||
if let Some(parent) = agent.parent {
|
||||
self.children.entry(parent).or_default().push(id);
|
||||
}
|
||||
|
||||
self.agents.push(agent);
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub(crate) fn is_direct_child(&self, parent: AgentId, child: AgentId) -> bool {
|
||||
self.children
|
||||
.get(&parent)
|
||||
.map(|kids| kids.contains(&child))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) fn descendants(&self, roots: &[AgentId]) -> Vec<AgentId> {
|
||||
let mut result = Vec::new();
|
||||
let mut stack: Vec<AgentId> = roots.to_vec();
|
||||
while let Some(id) = stack.pop() {
|
||||
result.push(id);
|
||||
if let Some(children) = self.children.get(&id) {
|
||||
for child in children {
|
||||
stack.push(*child);
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub(crate) fn next_sub_id(&mut self, agent: AgentId) -> String {
|
||||
let sub_id = format!("collab-agent-{}-{}", agent.0, self.next_sub_id);
|
||||
self.next_sub_id += 1;
|
||||
sub_id
|
||||
}
|
||||
|
||||
fn index_for(&self, id: AgentId) -> Option<usize> {
|
||||
if id.0 < 0 {
|
||||
return None;
|
||||
}
|
||||
let index = id.0 as usize;
|
||||
if index < self.agents.len() {
|
||||
Some(index)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn register_sub_id(&mut self, agent: AgentId, sub_id: String) {
|
||||
self.sub_ids.insert(sub_id, agent);
|
||||
}
|
||||
|
||||
pub(crate) fn agent_for_sub_id(&self, sub_id: &str) -> Option<AgentId> {
|
||||
self.sub_ids.get(sub_id).copied()
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,14 @@
|
||||
mod collaboration;
|
||||
mod service;
|
||||
mod session;
|
||||
mod turn;
|
||||
|
||||
pub(crate) use collaboration::AgentId;
|
||||
pub(crate) use collaboration::AgentLifecycleState;
|
||||
pub(crate) use collaboration::AgentState;
|
||||
pub(crate) use collaboration::CollaborationLimits;
|
||||
pub(crate) use collaboration::CollaborationState;
|
||||
pub(crate) use collaboration::ContextStrategy;
|
||||
pub(crate) use service::SessionServices;
|
||||
pub(crate) use session::SessionState;
|
||||
pub(crate) use turn::ActiveTurn;
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::codex::Session;
|
||||
use crate::codex::TurnContext;
|
||||
use crate::codex::run_collaboration_turn;
|
||||
use crate::response_processing::process_items;
|
||||
use crate::state::AgentId;
|
||||
use crate::state::AgentLifecycleState;
|
||||
use crate::state::TaskKind;
|
||||
use crate::tools::context::SharedTurnDiffTracker;
|
||||
use crate::turn_diff_tracker::TurnDiffTracker;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
|
||||
use super::SessionTask;
|
||||
use super::SessionTaskContext;
|
||||
|
||||
/// An async supervisor that drives child collaboration agents.
|
||||
///
|
||||
/// Each agent gets a runner task that executes turns on demand and emits
|
||||
/// `AgentRunResult` events on a broadcast channel. History swaps are serialized
|
||||
/// via a shared lock so runners can progress while the main agent continues.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CollaborationSupervisor {
|
||||
tx: mpsc::Sender<SupervisorCommand>,
|
||||
events: broadcast::Sender<AgentRunResult>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct AgentRunResult {
|
||||
pub(crate) agent: AgentId,
|
||||
pub(crate) delta_tokens: i32,
|
||||
pub(crate) status: AgentLifecycleState,
|
||||
pub(crate) last_message: Option<String>,
|
||||
pub(crate) sub_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum AgentCommand {
|
||||
Run { max_duration: i32 },
|
||||
Close,
|
||||
}
|
||||
|
||||
enum SupervisorCommand {
|
||||
RunAgents {
|
||||
targets: Vec<AgentId>,
|
||||
max_duration: i32,
|
||||
},
|
||||
CloseAgents {
|
||||
targets: Vec<AgentId>,
|
||||
},
|
||||
}
|
||||
|
||||
impl CollaborationSupervisor {
|
||||
pub(crate) fn spawn(session: Arc<Session>) -> Self {
|
||||
let (tx, mut rx) = mpsc::channel::<SupervisorCommand>(8);
|
||||
let (events, _rx) = broadcast::channel::<AgentRunResult>(64);
|
||||
let mut runners: HashMap<AgentId, mpsc::Sender<AgentCommand>> = HashMap::new();
|
||||
let events_tx = events.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(cmd) = rx.recv().await {
|
||||
match cmd {
|
||||
SupervisorCommand::RunAgents {
|
||||
targets,
|
||||
max_duration,
|
||||
} => {
|
||||
for agent in &targets {
|
||||
ensure_runner(
|
||||
*agent,
|
||||
&mut runners,
|
||||
Arc::clone(&session),
|
||||
events_tx.clone(),
|
||||
);
|
||||
}
|
||||
run_agents(&runners, targets, max_duration);
|
||||
}
|
||||
SupervisorCommand::CloseAgents { targets } => {
|
||||
for agent in targets {
|
||||
if let Some(tx) = runners.remove(&agent) {
|
||||
let _ = tx.send(AgentCommand::Close).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self { tx, events }
|
||||
}
|
||||
|
||||
pub(crate) fn subscribe(&self) -> broadcast::Receiver<AgentRunResult> {
|
||||
self.events.subscribe()
|
||||
}
|
||||
|
||||
pub(crate) async fn start_agents(
|
||||
&self,
|
||||
targets: Vec<AgentId>,
|
||||
max_duration: i32,
|
||||
) -> Result<(), String> {
|
||||
let cmd = SupervisorCommand::RunAgents {
|
||||
targets,
|
||||
max_duration,
|
||||
};
|
||||
self.tx
|
||||
.send(cmd)
|
||||
.await
|
||||
.map_err(|err| format!("collaboration supervisor unavailable: {err}"))
|
||||
}
|
||||
|
||||
pub(crate) async fn kick_agent(&self, agent: AgentId) {
|
||||
let tx = self.tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let _ = tx
|
||||
.send(SupervisorCommand::RunAgents {
|
||||
targets: vec![agent],
|
||||
max_duration: i32::MAX,
|
||||
})
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) async fn close_agents(&self, targets: Vec<AgentId>) {
|
||||
let _ = self
|
||||
.tx
|
||||
.send(SupervisorCommand::CloseAgents { targets })
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_runner(
|
||||
agent: AgentId,
|
||||
runners: &mut HashMap<AgentId, mpsc::Sender<AgentCommand>>,
|
||||
session: Arc<Session>,
|
||||
events: broadcast::Sender<AgentRunResult>,
|
||||
) {
|
||||
if runners.contains_key(&agent) {
|
||||
return;
|
||||
}
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<AgentCommand>(4);
|
||||
runners.insert(agent, tx);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut pending_run = true;
|
||||
let mut next_budget = i32::MAX;
|
||||
loop {
|
||||
if !pending_run {
|
||||
match rx.recv().await {
|
||||
Some(AgentCommand::Run { max_duration }) => {
|
||||
pending_run = true;
|
||||
next_budget = max_duration;
|
||||
}
|
||||
Some(AgentCommand::Close) | None => break,
|
||||
}
|
||||
}
|
||||
|
||||
if !pending_run {
|
||||
continue;
|
||||
}
|
||||
|
||||
let budget = next_budget;
|
||||
pending_run = false;
|
||||
next_budget = i32::MAX;
|
||||
|
||||
match run_agent_turns(Arc::clone(&session), agent, budget).await {
|
||||
Ok((results, keep_running)) => {
|
||||
for result in results {
|
||||
let _ = events.send(result);
|
||||
}
|
||||
pending_run = keep_running;
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = events.send(AgentRunResult {
|
||||
agent,
|
||||
delta_tokens: 0,
|
||||
status: AgentLifecycleState::Error { error: err },
|
||||
last_message: None,
|
||||
sub_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn run_agents(
|
||||
runners: &HashMap<AgentId, mpsc::Sender<AgentCommand>>,
|
||||
targets: Vec<AgentId>,
|
||||
max_duration: i32,
|
||||
) {
|
||||
for target in targets {
|
||||
if let Some(tx) = runners.get(&target) {
|
||||
let _ = tx.try_send(AgentCommand::Run { max_duration });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_agent_turns(
|
||||
session: Arc<Session>,
|
||||
target: AgentId,
|
||||
max_duration: i32,
|
||||
) -> Result<(Vec<AgentRunResult>, bool), String> {
|
||||
let mut remaining_budget = max_duration;
|
||||
let mut results = Vec::new();
|
||||
let mut keep_running = false;
|
||||
|
||||
while remaining_budget > 0 {
|
||||
let agent_snapshot = {
|
||||
let collab = session.collaboration_state().lock().await;
|
||||
collab.agent(target).cloned()
|
||||
};
|
||||
let Some(agent_snapshot) = agent_snapshot else {
|
||||
return Err(format!("unknown agent {}", target.0));
|
||||
};
|
||||
if matches!(
|
||||
agent_snapshot.status,
|
||||
AgentLifecycleState::Closed
|
||||
| AgentLifecycleState::Exhausted
|
||||
| AgentLifecycleState::Error { .. }
|
||||
) {
|
||||
results.push(AgentRunResult {
|
||||
agent: target,
|
||||
delta_tokens: 0,
|
||||
status: agent_snapshot.status,
|
||||
last_message: None,
|
||||
sub_id: None,
|
||||
});
|
||||
break;
|
||||
}
|
||||
{
|
||||
let mut collab = session.collaboration_state().lock().await;
|
||||
if let Some(agent) = collab.agent_mut(target) {
|
||||
agent.status = AgentLifecycleState::Running;
|
||||
}
|
||||
}
|
||||
|
||||
let mut agent_history = agent_snapshot.history.clone();
|
||||
let sub_id = {
|
||||
let mut collab = session.collaboration_state().lock().await;
|
||||
|
||||
collab.next_sub_id(target)
|
||||
};
|
||||
let turn_context =
|
||||
Arc::new(session.make_collaboration_turn_context(&agent_snapshot, sub_id.clone()));
|
||||
session.register_sub_id(target, sub_id.clone()).await;
|
||||
let tracker: SharedTurnDiffTracker =
|
||||
Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
|
||||
let mut agent_status = AgentLifecycleState::Running;
|
||||
let mut last_message: Option<String> = None;
|
||||
|
||||
let before_tokens = agent_history.get_total_token_usage();
|
||||
let run_result = run_collaboration_turn(
|
||||
Arc::clone(&session),
|
||||
Arc::clone(&turn_context),
|
||||
tracker,
|
||||
agent_history.get_history(),
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let (delta_tokens, continue_running) = match run_result {
|
||||
Ok(processed_items) => {
|
||||
let (responses, items_to_record) =
|
||||
process_items(processed_items, session.as_ref(), turn_context.as_ref()).await;
|
||||
let new_history = session.clone_history_for_agent(target).await;
|
||||
let after_tokens = new_history.get_total_token_usage();
|
||||
let delta_tokens = after_tokens
|
||||
.saturating_sub(before_tokens)
|
||||
.clamp(0, i32::MAX as i64) as i32;
|
||||
let last = crate::codex::get_last_assistant_message_from_turn(&items_to_record);
|
||||
{
|
||||
let mut collab = session.collaboration_state().lock().await;
|
||||
if let Some(agent) = collab.agent_mut(target) {
|
||||
agent_status = AgentLifecycleState::Idle {
|
||||
last_agent_message: last.clone(),
|
||||
};
|
||||
agent.status = agent_status.clone();
|
||||
agent.history = new_history.clone();
|
||||
}
|
||||
}
|
||||
last_message = last;
|
||||
(delta_tokens, !responses.is_empty())
|
||||
}
|
||||
Err(err) => {
|
||||
{
|
||||
let mut collab = session.collaboration_state().lock().await;
|
||||
if let Some(agent) = collab.agent_mut(target) {
|
||||
agent_status = AgentLifecycleState::Error {
|
||||
error: err.to_string(),
|
||||
};
|
||||
agent.status = agent_status.clone();
|
||||
}
|
||||
}
|
||||
(0, false)
|
||||
}
|
||||
};
|
||||
|
||||
remaining_budget = remaining_budget.saturating_sub(delta_tokens);
|
||||
|
||||
let final_status = {
|
||||
let collab = session.collaboration_state().lock().await;
|
||||
collab.agent(target).map(|a| a.status.clone())
|
||||
}
|
||||
.unwrap_or(agent_status.clone());
|
||||
|
||||
results.push(AgentRunResult {
|
||||
agent: target,
|
||||
delta_tokens: delta_tokens.clamp(0, i32::MAX),
|
||||
status: final_status,
|
||||
last_message,
|
||||
sub_id: Some(sub_id),
|
||||
});
|
||||
|
||||
keep_running |= continue_running;
|
||||
if !continue_running || remaining_budget <= 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok((results, keep_running))
|
||||
}
|
||||
|
||||
/// Collaboration task wrapper.
|
||||
#[allow(dead_code)]
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub(crate) struct CollaborationTask;
|
||||
|
||||
#[async_trait]
|
||||
impl SessionTask for CollaborationTask {
|
||||
fn kind(&self) -> TaskKind {
|
||||
TaskKind::Regular
|
||||
}
|
||||
|
||||
async fn run(
|
||||
self: Arc<Self>,
|
||||
session: Arc<SessionTaskContext>,
|
||||
ctx: Arc<TurnContext>,
|
||||
input: Vec<UserInput>,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> Option<String> {
|
||||
let sess = session.clone_session();
|
||||
crate::codex::run_task(sess, ctx, input, cancellation_token).await
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod collaboration;
|
||||
mod compact;
|
||||
mod ghost_snapshot;
|
||||
mod regular;
|
||||
@@ -28,9 +29,9 @@ use crate::state::RunningTask;
|
||||
use crate::state::TaskKind;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
|
||||
pub(crate) use collaboration::CollaborationSupervisor;
|
||||
pub(crate) use compact::CompactTask;
|
||||
pub(crate) use ghost_snapshot::GhostSnapshotTask;
|
||||
pub(crate) use regular::RegularTask;
|
||||
pub(crate) use review::ReviewTask;
|
||||
pub(crate) use undo::UndoTask;
|
||||
pub(crate) use user_shell::UserShellCommandTask;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,5 @@
|
||||
pub mod apply_patch;
|
||||
pub mod collaboration;
|
||||
mod grep_files;
|
||||
mod list_dir;
|
||||
mod mcp;
|
||||
@@ -13,6 +14,7 @@ mod view_image;
|
||||
pub use plan::PLAN_TOOL;
|
||||
|
||||
pub use apply_patch::ApplyPatchHandler;
|
||||
pub use collaboration::CollaborationHandler;
|
||||
pub use grep_files::GrepFilesHandler;
|
||||
pub use list_dir::ListDirHandler;
|
||||
pub use mcp::McpHandler;
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
|
||||
@@ -34,8 +34,7 @@ impl ToolRouter {
|
||||
config: &ToolsConfig,
|
||||
mcp_tools: Option<HashMap<String, mcp_types::Tool>>,
|
||||
) -> Self {
|
||||
let builder = build_specs(config, mcp_tools);
|
||||
let (specs, registry) = builder.build();
|
||||
let (specs, registry) = build_specs(config, mcp_tools).build();
|
||||
|
||||
Self { registry, specs }
|
||||
}
|
||||
@@ -136,6 +135,27 @@ impl ToolRouter {
|
||||
tracker: SharedTurnDiffTracker,
|
||||
call: ToolCall,
|
||||
) -> Result<ResponseInputItem, FunctionCallError> {
|
||||
let session_configuration = session.current_session_configuration().await;
|
||||
let session_history = session.clone_history().await;
|
||||
let mut caller_agent = turn.collaboration_agent();
|
||||
{
|
||||
let mut collab = session.collaboration_state().lock().await;
|
||||
collab.ensure_root_agent(&session_configuration, &session_history);
|
||||
if let Some(mapped) = collab.agent_for_sub_id(&turn.sub_id) {
|
||||
caller_agent = mapped;
|
||||
}
|
||||
if collab.agent(caller_agent).is_none() {
|
||||
return Ok(Self::failure_response(
|
||||
call.call_id,
|
||||
matches!(call.payload, ToolPayload::Custom { .. }),
|
||||
FunctionCallError::RespondToModel(format!(
|
||||
"unknown collaboration agent {}",
|
||||
caller_agent.0
|
||||
)),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let ToolCall {
|
||||
tool_name,
|
||||
call_id,
|
||||
@@ -187,3 +207,21 @@ impl ToolRouter {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::codex::make_session_and_context_with_rx;
|
||||
|
||||
#[test]
|
||||
fn builds_specs_without_filtering() {
|
||||
let (_session, turn, _rx) = make_session_and_context_with_rx();
|
||||
let router = ToolRouter::from_config(&turn.tools_config, None);
|
||||
let names: Vec<String> = router
|
||||
.specs()
|
||||
.into_iter()
|
||||
.map(|spec| spec.name().to_string())
|
||||
.collect();
|
||||
assert!(names.contains(&"update_plan".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,11 @@ use crate::tools::handlers::PLAN_TOOL;
|
||||
use crate::tools::handlers::apply_patch::ApplyPatchToolType;
|
||||
use crate::tools::handlers::apply_patch::create_apply_patch_freeform_tool;
|
||||
use crate::tools::handlers::apply_patch::create_apply_patch_json_tool;
|
||||
use crate::tools::handlers::collaboration::create_collaboration_close_tool;
|
||||
use crate::tools::handlers::collaboration::create_collaboration_get_state_tool;
|
||||
use crate::tools::handlers::collaboration::create_collaboration_init_agent_tool;
|
||||
use crate::tools::handlers::collaboration::create_collaboration_send_tool;
|
||||
use crate::tools::handlers::collaboration::create_collaboration_wait_tool;
|
||||
use crate::tools::registry::ToolRegistryBuilder;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
@@ -974,6 +979,7 @@ pub(crate) fn build_specs(
|
||||
mcp_tools: Option<HashMap<String, mcp_types::Tool>>,
|
||||
) -> ToolRegistryBuilder {
|
||||
use crate::tools::handlers::ApplyPatchHandler;
|
||||
use crate::tools::handlers::CollaborationHandler;
|
||||
use crate::tools::handlers::GrepFilesHandler;
|
||||
use crate::tools::handlers::ListDirHandler;
|
||||
use crate::tools::handlers::McpHandler;
|
||||
@@ -994,6 +1000,7 @@ pub(crate) fn build_specs(
|
||||
let plan_handler = Arc::new(PlanHandler);
|
||||
let apply_patch_handler = Arc::new(ApplyPatchHandler);
|
||||
let view_image_handler = Arc::new(ViewImageHandler);
|
||||
let collaboration_handler = Arc::new(CollaborationHandler);
|
||||
let mcp_handler = Arc::new(McpHandler);
|
||||
let mcp_resource_handler = Arc::new(McpResourceHandler);
|
||||
let shell_command_handler = Arc::new(ShellCommandHandler);
|
||||
@@ -1037,6 +1044,23 @@ pub(crate) fn build_specs(
|
||||
builder.push_spec(PLAN_TOOL.clone());
|
||||
builder.register_handler("update_plan", plan_handler);
|
||||
|
||||
if config
|
||||
.experimental_supported_tools
|
||||
.contains(&"collaboration".to_string())
|
||||
{
|
||||
builder.push_spec(create_collaboration_init_agent_tool());
|
||||
builder.push_spec(create_collaboration_send_tool());
|
||||
builder.push_spec(create_collaboration_wait_tool());
|
||||
builder.push_spec(create_collaboration_get_state_tool());
|
||||
builder.push_spec(create_collaboration_close_tool());
|
||||
|
||||
builder.register_handler("collaboration_init_agent", collaboration_handler.clone());
|
||||
builder.register_handler("collaboration_send", collaboration_handler.clone());
|
||||
builder.register_handler("collaboration_wait", collaboration_handler.clone());
|
||||
builder.register_handler("collaboration_get_state", collaboration_handler.clone());
|
||||
builder.register_handler("collaboration_close", collaboration_handler);
|
||||
}
|
||||
|
||||
if let Some(apply_patch_tool_type) = &config.apply_patch_tool_type {
|
||||
match apply_patch_tool_type {
|
||||
ApplyPatchToolType::Freeform => {
|
||||
@@ -1251,6 +1275,10 @@ mod tests {
|
||||
create_read_mcp_resource_tool(),
|
||||
PLAN_TOOL.clone(),
|
||||
create_apply_patch_freeform_tool(),
|
||||
create_grep_files_tool(),
|
||||
create_read_file_tool(),
|
||||
create_list_dir_tool(),
|
||||
create_test_sync_tool(),
|
||||
ToolSpec::WebSearch {},
|
||||
create_view_image_tool(),
|
||||
] {
|
||||
@@ -1296,6 +1324,10 @@ mod tests {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"view_image",
|
||||
],
|
||||
);
|
||||
@@ -1313,6 +1345,10 @@ mod tests {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"view_image",
|
||||
],
|
||||
);
|
||||
@@ -1333,6 +1369,10 @@ mod tests {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"web_search",
|
||||
"view_image",
|
||||
],
|
||||
@@ -1354,12 +1394,42 @@ mod tests {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"web_search",
|
||||
"view_image",
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_specs_gpt51_codex_max_default() {
|
||||
assert_model_tools(
|
||||
"gpt-5.1-codex-max",
|
||||
&Features::with_defaults(),
|
||||
&[
|
||||
"shell_command",
|
||||
"list_mcp_resources",
|
||||
"list_mcp_resource_templates",
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"collaboration_init_agent",
|
||||
"collaboration_send",
|
||||
"collaboration_wait",
|
||||
"collaboration_get_state",
|
||||
"collaboration_close",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"view_image",
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_codex_mini_defaults() {
|
||||
assert_model_tools(
|
||||
@@ -1388,6 +1458,10 @@ mod tests {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"view_image",
|
||||
],
|
||||
);
|
||||
|
||||
@@ -73,6 +73,10 @@ async fn model_selects_expected_tools() {
|
||||
"read_mcp_resource".to_string(),
|
||||
"update_plan".to_string(),
|
||||
"apply_patch".to_string(),
|
||||
"grep_files".to_string(),
|
||||
"read_file".to_string(),
|
||||
"list_dir".to_string(),
|
||||
"test_sync_tool".to_string(),
|
||||
"view_image".to_string()
|
||||
],
|
||||
"gpt-5-codex should expose the apply_patch tool",
|
||||
@@ -88,6 +92,10 @@ async fn model_selects_expected_tools() {
|
||||
"read_mcp_resource".to_string(),
|
||||
"update_plan".to_string(),
|
||||
"apply_patch".to_string(),
|
||||
"grep_files".to_string(),
|
||||
"read_file".to_string(),
|
||||
"list_dir".to_string(),
|
||||
"test_sync_tool".to_string(),
|
||||
"view_image".to_string()
|
||||
],
|
||||
"gpt-5.1-codex should expose the apply_patch tool",
|
||||
|
||||
@@ -159,6 +159,10 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> {
|
||||
"read_mcp_resource",
|
||||
"update_plan",
|
||||
"apply_patch",
|
||||
"grep_files",
|
||||
"read_file",
|
||||
"list_dir",
|
||||
"test_sync_tool",
|
||||
"view_image",
|
||||
];
|
||||
let body0 = req1.single_request().body_json();
|
||||
|
||||
Reference in New Issue
Block a user