llm: add dynamic tool loading (DTL) — Kimi system-tools + Anthropic tool-reference
Nightly Build / build (push) Successful in 6m51s
Nightly Build / build (push) Successful in 6m51s
Replace the old session_mcp_grants/stack_mcp_grants table pair with
a single activated_tools table that anchors each activation at the
assistant message_id that triggered it. The durable write moves from
the activate_tools tool itself to the round loop (handle_tool_call),
which has the message_id the DTL serializer positions injected tool
blocks against.
Introduce DtlMode (None / AnthropicToolReference / KimiSystemTools),
resolved per model from capabilities (opt-in via tool_search)
combined with the provider's dtl_format(). The message builder inserts
Kimi system {tools} blocks at the activation position, or emits
Anthropic tool_reference markers on the tool result. The tool-def
surface (all_tool_defs) switches shape: Anthropic declares everything
deferred; Kimi omits activated tools from the top-level array (system
takes over); None keeps the old grant-set logic.
Anthropic client: accept structured system arrays (cache_control on
the static block when DTL is active), carry defer_loading through
conversion, emit tool_reference blocks on result messages. Prompt
caching enabled exactly when DTL is active (anthropic provider).
MCP server list in the prompt is now a static catalogue (not split
Available/Active) — the split invalidated the cache on every activation.
Groundwork for providers.yaml dtl: key; Moonshot/Kimi providers wired
with kimi_system_tools and the k3* enrich now adds tool_search.
Compactor re-anchors activations whose message was compacted away.
This commit is contained in:
@@ -183,6 +183,15 @@ pub trait ApiProvider: Send + Sync {
|
||||
None
|
||||
}
|
||||
|
||||
/// The dynamic-tool-loading (DTL) serialization format this provider's models
|
||||
/// speak, e.g. `"anthropic_tool_reference"` or `"kimi_system_tools"`. Applied
|
||||
/// only to a model that opts in via the `tool_search` capability. `None` = no
|
||||
/// DTL (activated tools ride in the top-level `tools`). Returned as a string so
|
||||
/// core-api needs no dependency on the engine's `DtlMode` — the caller parses it.
|
||||
fn dtl_format(&self) -> Option<&str> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn llm_model_info(
|
||||
&self,
|
||||
_record: &LlmProviderRecord,
|
||||
|
||||
@@ -66,19 +66,36 @@ impl AnthropicClient {
|
||||
/// Converts OpenAI-format tool definitions to Anthropic format.
|
||||
/// OpenAI: { "type": "function", "function": { "name", "description", "parameters" } }
|
||||
/// Anthropic: { "name", "description", "input_schema" }
|
||||
///
|
||||
/// DTL (tool search): a top-level `defer_loading: true` on the OpenAI tool
|
||||
/// object is carried through to Anthropic's native `defer_loading` field. When
|
||||
/// any tool is deferred, the cache breakpoint is placed on the last
|
||||
/// **non-deferred** tool — a deferred tool cannot also carry `cache_control`
|
||||
/// (the API 400s), and at least one tool must stay non-deferred anyway.
|
||||
fn convert_tools(tools: &[Value]) -> Vec<Value> {
|
||||
tools
|
||||
let has_deferred = tools.iter().any(|t| t["defer_loading"].as_bool() == Some(true));
|
||||
let mut out: Vec<Value> = tools
|
||||
.iter()
|
||||
.filter_map(|t| {
|
||||
let func = &t["function"];
|
||||
let name = func["name"].as_str()?;
|
||||
Some(json!({
|
||||
let mut tool = json!({
|
||||
"name": name,
|
||||
"description": func["description"].as_str().unwrap_or(""),
|
||||
"input_schema": func["parameters"],
|
||||
}))
|
||||
});
|
||||
if t["defer_loading"].as_bool() == Some(true) {
|
||||
tool["defer_loading"] = json!(true);
|
||||
}
|
||||
Some(tool)
|
||||
})
|
||||
.collect()
|
||||
.collect();
|
||||
if has_deferred {
|
||||
if let Some(t) = out.iter_mut().rev().find(|t| t["defer_loading"].as_bool() != Some(true)) {
|
||||
t["cache_control"] = json!({ "type": "ephemeral" });
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Converts OpenAI-format message array to Anthropic format.
|
||||
@@ -145,10 +162,25 @@ impl AnthropicClient {
|
||||
let mut results: Vec<Value> = Vec::new();
|
||||
while i < messages.len() && messages[i]["role"].as_str() == Some("tool") {
|
||||
let tm = &messages[i];
|
||||
// DTL (custom tool search): a tool result carrying
|
||||
// `_tool_references` (set by the message builder on an
|
||||
// `activate_tools` result in AnthropicToolReference mode) becomes a
|
||||
// `content` array of `tool_reference` blocks, which the API expands
|
||||
// into the deferred tools' full definitions. Empty/absent → the
|
||||
// normal text result.
|
||||
let content: Value = match tm["_tool_references"].as_array() {
|
||||
Some(refs) if !refs.is_empty() => Value::Array(
|
||||
refs.iter()
|
||||
.filter_map(|r| r.as_str())
|
||||
.map(|name| json!({ "type": "tool_reference", "tool_name": name }))
|
||||
.collect(),
|
||||
),
|
||||
_ => Value::String(tm["content"].as_str().unwrap_or("").to_string()),
|
||||
};
|
||||
results.push(json!({
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tm["tool_call_id"].as_str().unwrap_or(""),
|
||||
"content": tm["content"].as_str().unwrap_or(""),
|
||||
"content": content,
|
||||
}));
|
||||
i += 1;
|
||||
}
|
||||
@@ -164,7 +196,7 @@ impl AnthropicClient {
|
||||
|
||||
/// Assembles the `/v1/messages` request body shared by the buffered and the
|
||||
/// streaming path (the caller adds `stream` on top).
|
||||
fn tools_body(&self, system: Option<String>, messages: Vec<Value>, tools: Vec<Value>, options: &ChatOptions) -> Value {
|
||||
fn tools_body(&self, system: Option<Value>, messages: Vec<Value>, tools: Vec<Value>, options: &ChatOptions) -> Value {
|
||||
let max_tokens = options.max_tokens.unwrap_or(4096);
|
||||
let mut body = json!({
|
||||
"model": options.model,
|
||||
@@ -173,22 +205,47 @@ impl AnthropicClient {
|
||||
"tools": tools,
|
||||
});
|
||||
|
||||
if let Some(sys) = system { body["system"] = sys.into(); }
|
||||
if let Some(sys) = system { body["system"] = sys; }
|
||||
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body);
|
||||
body
|
||||
}
|
||||
|
||||
/// Collects ALL system-role messages (main prompt, mid-conversation
|
||||
/// summary, tail_reminder) into a single `system:` string. The Anthropic
|
||||
/// API only accepts a single system parameter.
|
||||
fn merged_system(messages: &[Value]) -> Option<String> {
|
||||
let parts: Vec<&str> = messages
|
||||
/// Collects ALL system-role messages (main prompt, mid-conversation summary,
|
||||
/// tail_reminder) into the single `system` parameter the Anthropic API accepts.
|
||||
///
|
||||
/// Returns a plain string in the common case. When any system message carries
|
||||
/// **structured** content (a text-block array, e.g. the static prompt tagged
|
||||
/// with `cache_control` when prompt caching is on), it returns the array form
|
||||
/// instead so the cache breakpoint survives into `system`. String-content
|
||||
/// messages become plain text blocks (no cache_control).
|
||||
fn merged_system(messages: &[Value]) -> Option<Value> {
|
||||
let sys: Vec<&Value> = messages
|
||||
.iter()
|
||||
.filter(|m| m["role"].as_str() == Some("system"))
|
||||
.filter_map(|m| m["content"].as_str())
|
||||
.collect();
|
||||
if parts.is_empty() { None } else { Some(parts.join("\n\n---\n\n")) }
|
||||
if sys.is_empty() { return None; }
|
||||
|
||||
if !sys.iter().any(|m| m["content"].is_array()) {
|
||||
let parts: Vec<&str> = sys.iter().filter_map(|m| m["content"].as_str()).collect();
|
||||
return if parts.is_empty() { None } else { Some(Value::String(parts.join("\n\n---\n\n"))) };
|
||||
}
|
||||
|
||||
let mut blocks: Vec<Value> = Vec::new();
|
||||
for m in &sys {
|
||||
match &m["content"] {
|
||||
Value::String(s) if !s.is_empty() => blocks.push(json!({ "type": "text", "text": s })),
|
||||
Value::Array(arr) => {
|
||||
for b in arr {
|
||||
if b["type"].as_str() == Some("text") {
|
||||
blocks.push(b.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if blocks.is_empty() { None } else { Some(Value::Array(blocks)) }
|
||||
}
|
||||
|
||||
fn url(&self) -> String {
|
||||
|
||||
@@ -452,7 +452,7 @@ impl ChatHub {
|
||||
/// The next LLM turn will start with no MCP servers activated.
|
||||
pub async fn reset_mcp(&self, source_id: &str) -> anyhow::Result<()> {
|
||||
let session_id = self.get_or_create_session(source_id, &self.default_agent).await?;
|
||||
crate::db::session_mcp_grants::revoke_all(&self.db, session_id).await?;
|
||||
crate::db::activated_tools::revoke_all_session(&self.db, session_id).await?;
|
||||
info!(source_id, session_id, "ChatHub: MCP grants reset");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -338,6 +338,17 @@ impl ContextCompactor {
|
||||
|
||||
let summary_id = chat_summaries::save(pool, stack_id, &summary_text, last_covered_id).await?;
|
||||
|
||||
// DTL: activations pinned to a message that was just compacted away would
|
||||
// otherwise lose their render position (the Kimi `system`+`tools` block).
|
||||
// Re-anchor them onto the first surviving message. Best-effort — a failure
|
||||
// only means the model may re-activate a tool after compaction.
|
||||
let first_surviving_id = messages[split].id;
|
||||
if let Err(e) = crate::db::activated_tools::reanchor_compacted(
|
||||
pool, stack_id, last_covered_id, first_surviving_id,
|
||||
).await {
|
||||
warn!(stack_id, error = %e, "compactor: failed to re-anchor DTL activations");
|
||||
}
|
||||
|
||||
info!(
|
||||
stack_id,
|
||||
summary_id,
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
//! Persisted tool-group activations (the effect of `activate_tools`).
|
||||
//!
|
||||
//! One row per activated group, anchored at the assistant `message_id` that
|
||||
//! triggered it. Replaces the old `session_mcp_grants` / `stack_mcp_grants`
|
||||
//! pair: `stack_id IS NULL` is a session-scoped activation (root agent), a
|
||||
//! non-NULL `stack_id` is a sub-agent-frame activation (deleted on frame exit).
|
||||
//!
|
||||
//! The activation is the durable **effect**; the tool *call* itself lives in
|
||||
//! `chat_llm_tools`. Keeping them separate means "which groups are active" is a
|
||||
//! direct query, not a parse of `activate_tools` call arguments.
|
||||
//!
|
||||
//! `kind`/`ref` normalise the activated group: `('builtin', 'config')` for the
|
||||
//! reserved built-in group, `('mcp', <server name>)` for an MCP server.
|
||||
|
||||
use anyhow::Result;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
/// One activation row, anchored at the message that triggered it.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Activation {
|
||||
/// The assistant `chat_history.id` whose tool call triggered this activation.
|
||||
pub message_id: i64,
|
||||
/// `'builtin'` (the reserved `config` group) or `'mcp'` (a server).
|
||||
pub kind: String,
|
||||
/// The group reference: `'config'`, or the MCP server name.
|
||||
pub ref_: String,
|
||||
}
|
||||
|
||||
/// Persist a tool-group activation. `stack_id = None` → session-scoped (root
|
||||
/// agent); `Some(id)` → stack-scoped (sub-agent frame). Idempotent via the
|
||||
/// `COALESCE(stack_id, -1)`-based unique index (INSERT OR IGNORE).
|
||||
pub async fn grant(
|
||||
pool: &SqlitePool,
|
||||
session_id: i64,
|
||||
stack_id: Option<i64>,
|
||||
message_id: i64,
|
||||
kind: &str,
|
||||
ref_: &str,
|
||||
) -> Result<()> {
|
||||
sqlx::query(
|
||||
"INSERT OR IGNORE INTO activated_tools (session_id, stack_id, message_id, kind, ref)
|
||||
VALUES (?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(stack_id)
|
||||
.bind(message_id)
|
||||
.bind(kind)
|
||||
.bind(ref_)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Session-scoped activated group refs (root agent). The in-memory grant set is
|
||||
/// seeded from this at config-build time. Replaces
|
||||
/// `session_mcp_grants::list_for_session`.
|
||||
pub async fn list_refs_session(pool: &SqlitePool, session_id: i64) -> Result<Vec<String>> {
|
||||
let rows = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT ref FROM activated_tools
|
||||
WHERE session_id = ? AND stack_id IS NULL
|
||||
ORDER BY id",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(|(r,)| r).collect())
|
||||
}
|
||||
|
||||
/// Stack-scoped activated group refs (sub-agent frame). Sub-agents do **not**
|
||||
/// inherit session-scoped grants — they start from their own frame only, exactly
|
||||
/// as with the old `stack_mcp_grants::list_for_stack`.
|
||||
pub async fn list_refs_stack(pool: &SqlitePool, stack_id: i64) -> Result<Vec<String>> {
|
||||
let rows = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT ref FROM activated_tools
|
||||
WHERE stack_id = ?
|
||||
ORDER BY id",
|
||||
)
|
||||
.bind(stack_id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(|(r,)| r).collect())
|
||||
}
|
||||
|
||||
/// Activations in effect for one scope, up to and including `upto_message_id`,
|
||||
/// ordered by the anchoring message. Used by the DTL serializer to place the
|
||||
/// injected tool block (Kimi) at the position where it was activated. The scope
|
||||
/// mirrors the in-memory grant set: root (`stack_id = None`) sees session-scoped
|
||||
/// activations; a sub-agent (`Some(id)`) sees only its own frame's.
|
||||
pub async fn list_active_at(
|
||||
pool: &SqlitePool,
|
||||
session_id: i64,
|
||||
stack_id: Option<i64>,
|
||||
upto_message_id: i64,
|
||||
) -> Result<Vec<Activation>> {
|
||||
let rows = match stack_id {
|
||||
None => {
|
||||
sqlx::query_as::<_, (i64, String, String)>(
|
||||
"SELECT message_id, kind, ref FROM activated_tools
|
||||
WHERE session_id = ? AND stack_id IS NULL AND message_id <= ?
|
||||
ORDER BY message_id, id",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(upto_message_id)
|
||||
.fetch_all(pool)
|
||||
.await?
|
||||
}
|
||||
Some(sid) => {
|
||||
sqlx::query_as::<_, (i64, String, String)>(
|
||||
"SELECT message_id, kind, ref FROM activated_tools
|
||||
WHERE stack_id = ? AND message_id <= ?
|
||||
ORDER BY message_id, id",
|
||||
)
|
||||
.bind(sid)
|
||||
.bind(upto_message_id)
|
||||
.fetch_all(pool)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|(message_id, kind, ref_)| Activation { message_id, kind, ref_ })
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Clear all session-scoped activations for a session (the `/resettools` path).
|
||||
/// Stack-scoped rows are ephemeral (removed on frame exit) and there are none
|
||||
/// between turns, so only the session scope needs clearing.
|
||||
pub async fn revoke_all_session(pool: &SqlitePool, session_id: i64) -> Result<()> {
|
||||
sqlx::query("DELETE FROM activated_tools WHERE session_id = ? AND stack_id IS NULL")
|
||||
.bind(session_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a stack frame's activations. Called when the frame terminates.
|
||||
pub async fn delete_for_stack(pool: &SqlitePool, stack_id: i64) -> Result<()> {
|
||||
sqlx::query("DELETE FROM activated_tools WHERE stack_id = ?")
|
||||
.bind(stack_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Re-anchor activations pinned to a message of `stack_id` that was just compacted
|
||||
/// (chat_history id ≤ `covers_up_to`) onto `new_anchor` (the first surviving
|
||||
/// message), so the DTL serializer still renders them after compaction instead of
|
||||
/// losing the injection point. Scoped to this stack's messages via a subquery —
|
||||
/// `message_id` is a global autoincrement, so a bare `<=` would also match other
|
||||
/// stacks' rows. No unique-index conflict: only `message_id` changes, and there is
|
||||
/// at most one row per `(session, stack, kind, ref)`.
|
||||
pub async fn reanchor_compacted(
|
||||
pool: &SqlitePool,
|
||||
stack_id: i64,
|
||||
covers_up_to: i64,
|
||||
new_anchor: i64,
|
||||
) -> Result<()> {
|
||||
sqlx::query(
|
||||
"UPDATE activated_tools SET message_id = ?
|
||||
WHERE message_id IN (
|
||||
SELECT id FROM chat_history
|
||||
WHERE session_stack_id = ? AND id <= ?
|
||||
)",
|
||||
)
|
||||
.bind(new_anchor)
|
||||
.bind(stack_id)
|
||||
.bind(covers_up_to)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod activated_tools;
|
||||
pub mod approval_rules;
|
||||
pub mod project_members;
|
||||
pub mod projects;
|
||||
@@ -26,10 +27,8 @@ pub mod role_capabilities;
|
||||
pub mod roles;
|
||||
pub mod scheduled_jobs;
|
||||
pub mod scratchpad;
|
||||
pub mod session_mcp_grants;
|
||||
pub mod shared_folders;
|
||||
pub mod sources;
|
||||
pub mod stack_mcp_grants;
|
||||
pub mod tool_permission_groups;
|
||||
pub mod users;
|
||||
|
||||
@@ -809,26 +808,40 @@ pub async fn create_owner_tables(pool: &SqlitePool) -> Result<()> {
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
// Tool-group activations — the durable **effect** of `activate_tools`. One
|
||||
// row per activated group, anchored at the assistant `message_id` that
|
||||
// triggered it. `stack_id IS NULL` = session-scoped (root agent); non-NULL =
|
||||
// sub-agent frame (removed on frame exit). `kind`/`ref`: ('builtin','config')
|
||||
// or ('mcp', <server>). All FKs are owner→owner.
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS session_mcp_grants (
|
||||
"CREATE TABLE IF NOT EXISTS activated_tools (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id INTEGER NOT NULL,
|
||||
mcp_name TEXT NOT NULL,
|
||||
granted_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE(session_id, mcp_name)
|
||||
session_id INTEGER NOT NULL REFERENCES chat_sessions(id),
|
||||
stack_id INTEGER REFERENCES chat_sessions_stack(id),
|
||||
message_id INTEGER NOT NULL REFERENCES chat_history(id),
|
||||
kind TEXT NOT NULL CHECK(kind IN ('builtin', 'mcp')),
|
||||
ref TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
// Dedup per scope. A plain UNIQUE(session_id, stack_id, kind, ref) would NOT
|
||||
// dedup session-scoped rows: SQLite treats two NULL `stack_id`s as distinct,
|
||||
// so INSERT OR IGNORE would pile up duplicates. COALESCE folds NULL to -1.
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS stack_mcp_grants (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
stack_id INTEGER NOT NULL,
|
||||
mcp_name TEXT NOT NULL,
|
||||
granted_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE(stack_id, mcp_name)
|
||||
)",
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS ux_activated_tools
|
||||
ON activated_tools(session_id, COALESCE(stack_id, -1), kind, ref)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_activated_tools_stack ON activated_tools(stack_id)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_activated_tools_msg ON activated_tools(message_id)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
@@ -1080,8 +1093,9 @@ mod tests {
|
||||
one("INSERT INTO chat_summaries (stack_id, content, covers_up_to_message_id) VALUES (1, 's', 1)")
|
||||
.await.unwrap();
|
||||
one("INSERT INTO session_scratchpad (session_id, key, value) VALUES (1, 'k', 'v')").await.unwrap();
|
||||
one("INSERT INTO session_mcp_grants (session_id, mcp_name) VALUES (1, 'm')").await.unwrap();
|
||||
one("INSERT INTO stack_mcp_grants (stack_id, mcp_name) VALUES (1, 'm')").await.unwrap();
|
||||
// Both activation scopes: session-scoped (stack_id NULL) + stack-scoped.
|
||||
one("INSERT INTO activated_tools (session_id, stack_id, message_id, kind, ref) VALUES (1, NULL, 1, 'mcp', 'm')").await.unwrap();
|
||||
one("INSERT INTO activated_tools (session_id, stack_id, message_id, kind, ref) VALUES (1, 1, 1, 'mcp', 'm')").await.unwrap();
|
||||
one("INSERT INTO scheduled_jobs (id, title, cron, prompt, session_id) VALUES (1, 't', '* * * * *', 'p', 1)")
|
||||
.await.unwrap();
|
||||
one("INSERT INTO job_runs (job_id, started_at, status) VALUES (1, 'now', 'completed')").await.unwrap();
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
/// Grant access to an MCP server for a session.
|
||||
/// Uses INSERT OR IGNORE so calling it multiple times is safe.
|
||||
pub async fn grant(pool: &SqlitePool, session_id: i64, mcp_name: &str) -> Result<()> {
|
||||
sqlx::query(
|
||||
"INSERT OR IGNORE INTO session_mcp_grants (session_id, mcp_name)
|
||||
VALUES (?, ?)"
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(mcp_name)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Revoke all MCP grants for a session.
|
||||
pub async fn revoke_all(pool: &SqlitePool, session_id: i64) -> Result<()> {
|
||||
sqlx::query("DELETE FROM session_mcp_grants WHERE session_id = ?")
|
||||
.bind(session_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the names of all MCP servers granted for this session.
|
||||
pub async fn list_for_session(pool: &SqlitePool, session_id: i64) -> Result<Vec<String>> {
|
||||
let rows = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT mcp_name FROM session_mcp_grants WHERE session_id = ? ORDER BY granted_at"
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(|(name,)| name).collect())
|
||||
}
|
||||
@@ -1,36 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
/// Persist an MCP grant scoped to a specific stack frame (sub-agent).
|
||||
/// Uses INSERT OR IGNORE so calling it multiple times is safe.
|
||||
pub async fn grant(pool: &SqlitePool, stack_id: i64, mcp_name: &str) -> Result<()> {
|
||||
sqlx::query(
|
||||
"INSERT OR IGNORE INTO stack_mcp_grants (stack_id, mcp_name)
|
||||
VALUES (?, ?)",
|
||||
)
|
||||
.bind(stack_id)
|
||||
.bind(mcp_name)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the names of all MCP servers granted for this stack frame.
|
||||
pub async fn list_for_stack(pool: &SqlitePool, stack_id: i64) -> Result<Vec<String>> {
|
||||
let rows = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT mcp_name FROM stack_mcp_grants WHERE stack_id = ? ORDER BY granted_at",
|
||||
)
|
||||
.bind(stack_id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(|(name,)| name).collect())
|
||||
}
|
||||
|
||||
/// Removes all MCP grants for a stack frame. Called when the frame terminates.
|
||||
pub async fn delete_for_stack(pool: &SqlitePool, stack_id: i64) -> Result<()> {
|
||||
sqlx::query("DELETE FROM stack_mcp_grants WHERE stack_id = ?")
|
||||
.bind(stack_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -531,6 +531,16 @@ fn build_entry(
|
||||
context_length: model.context_length,
|
||||
prompt_cache,
|
||||
capabilities: model.capabilities.clone(),
|
||||
// DTL is opt-in per model (the `tool_search` capability); the wire *format*
|
||||
// comes from the model's provider (native, or `providers.yaml`) — no
|
||||
// hardcoded model list.
|
||||
dtl: if model.capabilities.iter().any(|c| c == "tool_search") {
|
||||
registry.get(&provider.provider)
|
||||
.and_then(|p| p.dtl_format().map(crate::llm::dtl_mode_from_format))
|
||||
.unwrap_or(crate::llm::DtlMode::None)
|
||||
} else {
|
||||
crate::llm::DtlMode::None
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -26,6 +26,50 @@ pub struct LlmEntry {
|
||||
/// Input capabilities of the resolved model (`vision`, `video`, …), from
|
||||
/// `llm_models.capabilities`. Drives multimodal attachment inlining.
|
||||
pub capabilities: Vec<String>,
|
||||
/// Dynamic-tool-loading serialization mode for this model (resolved from
|
||||
/// `capabilities` + provider type). Selects how a session's *activated* tools
|
||||
/// are put on the wire so that activating one does not invalidate the
|
||||
/// provider's prompt-cache prefix.
|
||||
pub dtl: DtlMode,
|
||||
}
|
||||
|
||||
/// Per-model dynamic-tool-loading (DTL) serialization mode. Resolved in
|
||||
/// `build_entry` from the model's provider (via [`dtl_mode_from_format`]) gated by
|
||||
/// the `tool_search` capability. It selects how a session's activated tools are serialized so
|
||||
/// that an `activate_tools` call does not break the provider's prompt-cache
|
||||
/// prefix. The persistence layer (`activated_tools`) is model-agnostic; this is
|
||||
/// the model-aware half that renders that state per provider.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum DtlMode {
|
||||
/// Today's behaviour: activated tools ride in the top-level `tools` array.
|
||||
/// Correct, but every activation invalidates the cache from that point on.
|
||||
/// The fallback for Ollama / LM Studio / generic OpenAI-compat providers.
|
||||
#[default]
|
||||
None,
|
||||
/// Anthropic Messages API: candidate tools are declared `defer_loading:true`
|
||||
/// and loaded via a custom client-side `tool_reference` expansion emitted at
|
||||
/// the `activate_tools` result (preserves the cache; no 5-result cap).
|
||||
AnthropicToolReference,
|
||||
/// Kimi K3 (OpenAI-compatible): activated tools are injected as `system`
|
||||
/// messages carrying a `tools` field, appended at the activation position so
|
||||
/// the prefix stays byte-identical (append-only).
|
||||
KimiSystemTools,
|
||||
}
|
||||
|
||||
/// Parses a provider-declared DTL format name — from a native provider
|
||||
/// (`AnthropicProvider::dtl_format`) or from `providers.yaml` (`dtl:` on a declared
|
||||
/// provider) — into a [`DtlMode`]. Unknown names → [`DtlMode::None`].
|
||||
///
|
||||
/// The *format* is a property of the provider (which wire its client speaks);
|
||||
/// whether a given model *uses* it is gated separately by the `tool_search`
|
||||
/// capability (see `build_entry`). So there is no hardcoded model list — enabling
|
||||
/// a new Kimi-compatible provider is a `providers.yaml` edit.
|
||||
pub fn dtl_mode_from_format(fmt: &str) -> DtlMode {
|
||||
match fmt {
|
||||
"anthropic_tool_reference" => DtlMode::AnthropicToolReference,
|
||||
"kimi_system_tools" => DtlMode::KimiSystemTools,
|
||||
_ => DtlMode::None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Provider ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -25,6 +25,12 @@ impl ApiProvider for AnthropicProvider {
|
||||
&[ServiceType::Llm]
|
||||
}
|
||||
|
||||
fn dtl_format(&self) -> Option<&str> {
|
||||
// Every Anthropic model that opts in (via the `tool_search` capability)
|
||||
// uses the custom client-side tool_reference format.
|
||||
Some("anthropic_tool_reference")
|
||||
}
|
||||
|
||||
async fn list_llm_models(&self, _record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
|
||||
Ok(None)
|
||||
}
|
||||
@@ -97,9 +103,16 @@ impl ApiProvider for AnthropicProvider {
|
||||
.with_context(|| format!("provider '{}': api_key required for anthropic", record.name))?;
|
||||
// Merge model extra_params + reasoning (thinking) into the request body.
|
||||
let extra = extra_with_reasoning(self, model);
|
||||
// Prompt caching is enabled exactly when this model runs dynamic tool
|
||||
// loading (the `tool_search` capability → custom tool_reference): the
|
||||
// deferred toolset keeps the tools prefix stable and the message builder
|
||||
// tags the static system block with cache_control, which the client
|
||||
// renders into the `system` array. Without DTL the native Anthropic path
|
||||
// stays uncached, as before.
|
||||
let prompt_cache = model.capabilities.iter().any(|c| c == "tool_search");
|
||||
Ok(BuiltLlmClient {
|
||||
client: Arc::new(AnthropicClient::with_extra_body(key, extra)),
|
||||
prompt_cache: false,
|
||||
prompt_cache,
|
||||
})
|
||||
})())
|
||||
}
|
||||
|
||||
@@ -59,6 +59,11 @@ struct ProviderSpec {
|
||||
fields: Vec<FieldSpec>,
|
||||
models: Option<ModelsSpec>,
|
||||
reasoning: Option<ReasoningSpec>,
|
||||
/// Dynamic-tool-loading wire format for this provider's models (e.g.
|
||||
/// `kimi_system_tools`). Applied only to a model with the `tool_search`
|
||||
/// capability. Absent → no DTL for this provider.
|
||||
#[serde(default)]
|
||||
dtl: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, serde::Deserialize)]
|
||||
@@ -506,6 +511,10 @@ impl ApiProvider for DeclaredProvider {
|
||||
LLM_ONLY
|
||||
}
|
||||
|
||||
fn dtl_format(&self) -> Option<&str> {
|
||||
self.spec.dtl.as_deref()
|
||||
}
|
||||
|
||||
async fn list_llm_models(
|
||||
&self,
|
||||
record: &LlmProviderRecord,
|
||||
|
||||
@@ -6,7 +6,7 @@ use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::info;
|
||||
|
||||
use crate::db::{chat_history, chat_llm_tools, chat_sessions_stack, scratchpad, stack_mcp_grants};
|
||||
use crate::db::{activated_tools, chat_history, chat_llm_tools, chat_sessions_stack, scratchpad};
|
||||
use crate::events::ServerEvent;
|
||||
|
||||
use super::{ChatSessionHandler, MAX_AGENT_DEPTH, TurnOutcome};
|
||||
@@ -108,8 +108,8 @@ impl ChatSessionHandler {
|
||||
// Sub-agents never inject live user input.
|
||||
let outcome = self.run_agent_turn(child.id, &child_config, token, tx, None).await;
|
||||
|
||||
if let Err(e) = stack_mcp_grants::delete_for_stack(pool, child.id).await {
|
||||
tracing::warn!(stack_id = child.id, error = %e, "dispatch_sub_agent: failed to delete stack MCP grants");
|
||||
if let Err(e) = activated_tools::delete_for_stack(pool, child.id).await {
|
||||
tracing::warn!(stack_id = child.id, error = %e, "dispatch_sub_agent: failed to delete stack activations");
|
||||
}
|
||||
|
||||
let parent_agent_id = parent_config.agent_id.clone();
|
||||
@@ -165,7 +165,7 @@ impl ChatSessionHandler {
|
||||
stack_id: i64,
|
||||
depth: i64,
|
||||
) -> anyhow::Result<AgentRunConfig> {
|
||||
let persisted_grants = stack_mcp_grants::list_for_stack(&self.db, stack_id)
|
||||
let persisted_grants = activated_tools::list_refs_stack(&self.db, stack_id)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let active_mcp_grants: Arc<RwLock<HashSet<String>>> =
|
||||
@@ -205,8 +205,6 @@ impl ChatSessionHandler {
|
||||
|
||||
{
|
||||
let activate_tool = crate::tools::activate_tools::ActivateTools {
|
||||
pool: Arc::clone(&self.db),
|
||||
session_id: self.session_id,
|
||||
stack_id: Some(stack_id),
|
||||
mcp: Arc::clone(&self.mcp),
|
||||
active_mcp_grants: Arc::clone(&active_mcp_grants),
|
||||
|
||||
@@ -140,7 +140,7 @@ impl ChatSessionHandler {
|
||||
// Load persisted session grants from DB (MCP server names and/or the reserved
|
||||
// `config` keyword), then inject `activate_tools` so the LLM can activate
|
||||
// additional groups on demand.
|
||||
let persisted = crate::db::session_mcp_grants::list_for_session(
|
||||
let persisted = crate::db::activated_tools::list_refs_session(
|
||||
&self.db, self.session_id,
|
||||
).await.unwrap_or_default();
|
||||
|
||||
@@ -148,14 +148,10 @@ impl ChatSessionHandler {
|
||||
Arc::new(RwLock::new(persisted.into_iter().collect()));
|
||||
|
||||
{
|
||||
let pool_clone = Arc::clone(&self.db);
|
||||
let session_id = self.session_id;
|
||||
let mcp_clone = Arc::clone(&self.mcp);
|
||||
let grants_clone = Arc::clone(&active_mcp_grants);
|
||||
|
||||
let activate_tool = crate::tools::activate_tools::ActivateTools {
|
||||
pool: pool_clone,
|
||||
session_id,
|
||||
stack_id: None,
|
||||
mcp: mcp_clone,
|
||||
active_mcp_grants: grants_clone,
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::sync::{Arc, RwLock};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm::DtlMode;
|
||||
use crate::mcp::McpProvider;
|
||||
use crate::tools::Tool;
|
||||
use crate::tools::tool_names as tn;
|
||||
@@ -56,10 +57,10 @@ pub struct AgentRunConfig {
|
||||
pub mcp: Arc<dyn McpProvider>,
|
||||
/// Set of MCP server names currently granted (activated) for this agent run.
|
||||
///
|
||||
/// - Root agents: pre-populated from `session_mcp_grants` DB at config-build time;
|
||||
/// updated in-place by `activate_tools`.
|
||||
/// - Root agents: pre-populated from the `activated_tools` table (session-scoped
|
||||
/// rows) at config-build time; updated in-place by `activate_tools`.
|
||||
/// - Sub-agents: starts empty; populated by `activate_tools` (stack-scoped, no
|
||||
/// session leak); deleted from DB when the stack frame terminates.
|
||||
/// session leak); the frame's rows are deleted when the stack frame terminates.
|
||||
///
|
||||
/// May also contain the reserved keyword `"config"`, which unlocks the built-in
|
||||
/// `Config`-category tools (`config_tool_defs`) rather than an MCP server.
|
||||
@@ -79,31 +80,39 @@ impl AgentRunConfig {
|
||||
///
|
||||
/// Dynamic groups are re-queried every call so that an `activate_tools` call in
|
||||
/// round N makes the tools visible in round N+1 without rebuilding the whole config.
|
||||
pub fn all_tool_defs(&self) -> Vec<Value> {
|
||||
pub fn all_tool_defs(&self, dtl: DtlMode) -> Vec<Value> {
|
||||
let mut defs = self.base_tool_defs.clone();
|
||||
|
||||
// Dynamic groups: read the currently-granted set (MCP server names + `config`).
|
||||
let granted: HashSet<String> = self.active_mcp_grants
|
||||
.read()
|
||||
.map(|g| g.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
// MCP servers: include tools for the granted server names.
|
||||
let servers: Vec<String> = granted.iter()
|
||||
.filter(|n| n.as_str() != crate::tools::tool_names::CONFIG_GROUP)
|
||||
.cloned()
|
||||
.collect();
|
||||
if !servers.is_empty() {
|
||||
defs.extend(
|
||||
self.mcp.tools_for(&servers)
|
||||
.iter()
|
||||
.map(|t| t.to_openai_definition()),
|
||||
);
|
||||
}
|
||||
|
||||
// `config` group: include the built-in Config-category tools on demand.
|
||||
if granted.contains(crate::tools::tool_names::CONFIG_GROUP) {
|
||||
defs.extend(self.config_tool_defs.iter().cloned());
|
||||
match dtl {
|
||||
// Anthropic custom tool_reference: declare EVERY accessible MCP tool +
|
||||
// the config group as `defer_loading:true`, on every turn. The toolset
|
||||
// is stable (cache-safe) and `activate_tools` loads the needed ones via
|
||||
// tool_reference. Deferred defs are excluded from the prompt prefix by
|
||||
// the API and cost nothing until referenced.
|
||||
DtlMode::AnthropicToolReference => {
|
||||
defs.extend(self.mcp.tools().iter().map(|t| deferred(t.to_openai_definition())));
|
||||
defs.extend(self.config_tool_defs.iter().cloned().map(deferred));
|
||||
}
|
||||
// Kimi K3: activated MCP/config tools are injected as `system` messages
|
||||
// by the message builder, so they are NOT in the top-level tools here.
|
||||
DtlMode::KimiSystemTools => {}
|
||||
// Today's behaviour: only the currently-granted MCP servers + config.
|
||||
DtlMode::None => {
|
||||
let granted: HashSet<String> = self.active_mcp_grants
|
||||
.read()
|
||||
.map(|g| g.clone())
|
||||
.unwrap_or_default();
|
||||
let servers: Vec<String> = granted.iter()
|
||||
.filter(|n| n.as_str() != crate::tools::tool_names::CONFIG_GROUP)
|
||||
.cloned()
|
||||
.collect();
|
||||
if !servers.is_empty() {
|
||||
defs.extend(self.mcp.tools_for(&servers).iter().map(|t| t.to_openai_definition()));
|
||||
}
|
||||
if granted.contains(crate::tools::tool_names::CONFIG_GROUP) {
|
||||
defs.extend(self.config_tool_defs.iter().cloned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
defs.extend(self.memory_tools.iter().map(|t| t.openai_definition()));
|
||||
@@ -167,3 +176,12 @@ impl AgentRunConfig {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tags an OpenAI tool definition as deferred (Anthropic tool search): the API
|
||||
/// keeps it out of the prompt prefix until `activate_tools` references it. The
|
||||
/// flag rides on the top-level tool object; `AnthropicClient::convert_tools`
|
||||
/// maps it to Anthropic's native `defer_loading` field.
|
||||
fn deferred(mut def: Value) -> Value {
|
||||
def["defer_loading"] = Value::Bool(true);
|
||||
def
|
||||
}
|
||||
|
||||
@@ -45,7 +45,6 @@ impl ChatSessionHandler {
|
||||
stack_id: i64,
|
||||
config: &AgentRunConfig,
|
||||
active_grants: &HashSet<String>,
|
||||
tool_defs: &[Value],
|
||||
req_scope: Option<&str>,
|
||||
req_strength: Option<LlmStrength>,
|
||||
cur_name: &mut String,
|
||||
@@ -57,6 +56,9 @@ impl ChatSessionHandler {
|
||||
let mut tried_this_round: Vec<String> = vec![cur_name.clone()];
|
||||
|
||||
loop {
|
||||
// Re-derive the tool defs for the model actually serving this attempt:
|
||||
// a fallback across DTL modes must re-shape (deferred candidates or not).
|
||||
let cur_tool_defs = config.all_tool_defs(cur_llm.dtl);
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let options = ChatOptions {
|
||||
model: cur_llm.model.clone(),
|
||||
@@ -72,8 +74,8 @@ impl ChatSessionHandler {
|
||||
// open directly — keyed on the model actually serving this attempt, so a
|
||||
// fallback to a text-only model drops the claim. `None` (no media
|
||||
// capability) leaves the shared defs untouched, avoiding a clone.
|
||||
let annotated = media_annotated_tools(tool_defs, &cur_llm.capabilities);
|
||||
let defs: &[Value] = annotated.as_deref().unwrap_or(tool_defs);
|
||||
let annotated = media_annotated_tools(&cur_tool_defs, &cur_llm.capabilities);
|
||||
let defs: &[Value] = annotated.as_deref().unwrap_or(&cur_tool_defs);
|
||||
|
||||
// Clone the Arc so the in-flight future does not borrow `cur_llm` across
|
||||
// the fallback reassignment below. On cancel we drop the future
|
||||
@@ -164,11 +166,13 @@ impl ChatSessionHandler {
|
||||
// settings (e.g. switching from OpenRouter/Anthropic to DeepSeek)
|
||||
// or different input capabilities (a non-vision fallback drops
|
||||
// inline media back to the textual path block).
|
||||
let activation_stack = if config.depth == 0 { None } else { Some(stack_id) };
|
||||
match self.build_openai_messages(
|
||||
&self.db, stack_id, &config.agent_id,
|
||||
config.extra_system.as_deref(), config.extra_system_dynamic.as_deref(),
|
||||
config.tail_reminder.as_deref(), active_grants,
|
||||
&config.system_substitutions, cur_llm.prompt_cache, &cur_llm.capabilities,
|
||||
cur_llm.dtl, &config.config_tool_defs, activation_stack,
|
||||
).await {
|
||||
Ok(m) => *messages = m,
|
||||
Err(e) => return RoundLlm::Failed(e),
|
||||
|
||||
@@ -117,8 +117,11 @@ impl ChatSessionHandler {
|
||||
// Messages are (re)built with the current model's prompt_cache flag.
|
||||
// On fallback within the same round `call_llm_round` rebuilds them again
|
||||
// if the replacement model has a different prompt_cache setting.
|
||||
let mut messages = self.build_openai_messages(pool, stack_id, &config.agent_id, config.extra_system.as_deref(), config.extra_system_dynamic.as_deref(), config.tail_reminder.as_deref(), &active_grants_snapshot, &config.system_substitutions, cur_llm.prompt_cache, &cur_llm.capabilities).await?;
|
||||
let tool_defs = config.all_tool_defs();
|
||||
// Activation scope for the DTL serializer: session-scoped for the root
|
||||
// agent (stack_id NULL), the frame itself for a sub-agent.
|
||||
let activation_stack = if config.depth == 0 { None } else { Some(stack_id) };
|
||||
let mut messages = self.build_openai_messages(pool, stack_id, &config.agent_id, config.extra_system.as_deref(), config.extra_system_dynamic.as_deref(), config.tail_reminder.as_deref(), &active_grants_snapshot, &config.system_substitutions, cur_llm.prompt_cache, &cur_llm.capabilities, cur_llm.dtl, &config.config_tool_defs, activation_stack).await?;
|
||||
let tool_defs = config.all_tool_defs(cur_llm.dtl);
|
||||
|
||||
// Record every tool actually offered to the LLM so the Security-groups
|
||||
// UI can list/gate dynamically-injected tools. Cheap no-op once each
|
||||
@@ -128,7 +131,7 @@ impl ChatSessionHandler {
|
||||
// One LLM call for this round, with automatic model fallback on
|
||||
// retriable errors. `cur_name`/`cur_llm`/`messages` are updated in place.
|
||||
let turn_result = match self.call_llm_round(
|
||||
stack_id, config, &active_grants_snapshot, &tool_defs,
|
||||
stack_id, config, &active_grants_snapshot,
|
||||
req_scope.as_deref(), req_strength,
|
||||
&mut cur_name, &mut cur_llm, &mut messages, token, &em,
|
||||
).await {
|
||||
@@ -274,6 +277,25 @@ impl ChatSessionHandler {
|
||||
DispatchResult::AbortPending => return Ok(CallFlow::End(TurnOutcome::Cancelled)),
|
||||
};
|
||||
|
||||
// Persist the durable effect of `activate_tools`, anchored at the assistant
|
||||
// `message_id` that triggered it (the anchor the DTL serializer positions
|
||||
// injected tool blocks against). The in-memory grant set was already updated
|
||||
// inside the tool; this records it across turns and restarts.
|
||||
if call.name == crate::tools::tool_names::ACTIVATE_TOOLS {
|
||||
if let Some(groups) = call.arguments.get("groups").and_then(|g| g.as_array()) {
|
||||
// Root (depth 0) → session-scoped (stack_id NULL); sub-agent → its frame.
|
||||
let anchor_stack = if config.depth == 0 { None } else { Some(stack_id) };
|
||||
for g in groups.iter().filter_map(|v| v.as_str()) {
|
||||
let kind = if g == crate::tools::tool_names::CONFIG_GROUP { "builtin" } else { "mcp" };
|
||||
if let Err(e) = crate::db::activated_tools::grant(
|
||||
pool, self.session_id, anchor_stack, message_id, kind, g,
|
||||
).await {
|
||||
tracing::warn!(session_id = self.session_id, group = g, error = %e, "activate_tools: failed to persist activation");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match self.record_tool_outcome(
|
||||
tool_call_id, &call.name, &call.arguments, outcome, preview, em, Some(all_tool_calls),
|
||||
).await? {
|
||||
|
||||
@@ -10,6 +10,7 @@ use core_api::user_fs::UserFs;
|
||||
use crate::compactor::{ContextCompactor, SUMMARY_PREFIX};
|
||||
use crate::config::DatetimeConfig;
|
||||
use crate::db::{chat_history, chat_llm_tools, chat_summaries};
|
||||
use crate::llm::DtlMode;
|
||||
use crate::mcp::McpProvider;
|
||||
use crate::tools::tool_names as tn;
|
||||
|
||||
@@ -108,6 +109,12 @@ impl MessageBuilder {
|
||||
// Input capabilities of the resolved model (`vision`, `video`, …) —
|
||||
// drives inline media for current-turn attachments.
|
||||
capabilities: &[String],
|
||||
// Dynamic-tool-loading mode for the resolved model, plus the config-group
|
||||
// tool defs (needed to resolve activated tools). `activation_stack` scopes
|
||||
// the activation read: `None` = session-scoped (root), `Some(id)` = frame.
|
||||
dtl: DtlMode,
|
||||
config_tool_defs: &[Value],
|
||||
activation_stack: Option<i64>,
|
||||
) -> anyhow::Result<Vec<Value>> {
|
||||
let pool = &*self.pool;
|
||||
|
||||
@@ -254,6 +261,15 @@ impl MessageBuilder {
|
||||
media_turn_start -= 1;
|
||||
}
|
||||
|
||||
// DTL: resolve the tools activated at each assistant message (empty unless a
|
||||
// DTL mode is active). Drives the Kimi `system`+`tools` injection and the
|
||||
// Anthropic `tool_reference` markers emitted in the history loop below.
|
||||
let activation_defs: HashMap<i64, Vec<Value>> = if matches!(dtl, DtlMode::None) {
|
||||
HashMap::new()
|
||||
} else {
|
||||
self.resolve_activation_defs(activation_stack, config_tool_defs).await
|
||||
};
|
||||
|
||||
for (idx, entry) in history.iter().enumerate() {
|
||||
let is_previous_turn = current_turn_boundary.is_some_and(|b| idx < b);
|
||||
|
||||
@@ -378,11 +394,29 @@ impl MessageBuilder {
|
||||
tc.arguments.as_deref(),
|
||||
);
|
||||
|
||||
out.push(json!({
|
||||
let mut tool_msg = json!({
|
||||
"role": "tool",
|
||||
"tool_call_id": format!("tc_{}", tc.id),
|
||||
"content": result_content,
|
||||
}));
|
||||
});
|
||||
// Anthropic DTL: an `activate_tools` result becomes a set of
|
||||
// `tool_reference`s (the activated groups' tool names) that the
|
||||
// client renders as tool_reference blocks and the API expands
|
||||
// into the deferred tools. `_tool_references` is a neutral
|
||||
// marker other clients ignore.
|
||||
if matches!(dtl, DtlMode::AnthropicToolReference)
|
||||
&& tc.name == tn::ACTIVATE_TOOLS
|
||||
&& let Some(adefs) = activation_defs.get(&entry.id)
|
||||
{
|
||||
let names: Vec<Value> = adefs.iter()
|
||||
.filter_map(|d| d["function"]["name"].as_str())
|
||||
.map(|n| Value::String(n.to_string()))
|
||||
.collect();
|
||||
if !names.is_empty() {
|
||||
tool_msg["_tool_references"] = Value::Array(names);
|
||||
}
|
||||
}
|
||||
out.push(tool_msg);
|
||||
}
|
||||
|
||||
// Media a tool produced this turn (e.g. read_file on an
|
||||
@@ -411,6 +445,16 @@ impl MessageBuilder {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Kimi K3 DTL: inject the tools activated at this assistant
|
||||
// message as a `system` message carrying a `tools` field, right
|
||||
// after its tool-result group (append-only → cache-safe).
|
||||
if matches!(dtl, DtlMode::KimiSystemTools)
|
||||
&& let Some(adefs) = activation_defs.get(&entry.id)
|
||||
&& !adefs.is_empty()
|
||||
{
|
||||
out.push(json!({ "role": "system", "tools": adefs }));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -585,7 +629,49 @@ impl MessageBuilder {
|
||||
))
|
||||
}
|
||||
|
||||
fn render_mcp_list(&self, active_mcp_grants: &HashSet<String>) -> String {
|
||||
/// Resolves the tools activated at each assistant message, for DTL
|
||||
/// serialization. Returns `message_id → deduped OpenAI tool defs`. MCP groups
|
||||
/// resolve to the server's live tool defs; the `config` group resolves to the
|
||||
/// passed-in config-category defs. Scope mirrors the in-memory grant set:
|
||||
/// `None` = session-scoped (root), `Some(id)` = the sub-agent frame.
|
||||
async fn resolve_activation_defs(
|
||||
&self,
|
||||
activation_stack: Option<i64>,
|
||||
config_tool_defs: &[Value],
|
||||
) -> HashMap<i64, Vec<Value>> {
|
||||
let activations = crate::db::activated_tools::list_active_at(
|
||||
&self.pool, self.session_id, activation_stack, i64::MAX,
|
||||
).await.unwrap_or_default();
|
||||
|
||||
let mut map: HashMap<i64, Vec<Value>> = HashMap::new();
|
||||
for a in activations {
|
||||
let resolved: Vec<Value> = match a.kind.as_str() {
|
||||
"mcp" => self.mcp.tools_for(&[a.ref_.clone()])
|
||||
.iter().map(|t| t.to_openai_definition()).collect(),
|
||||
"builtin" if a.ref_ == tn::CONFIG_GROUP => config_tool_defs.to_vec(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
let entry = map.entry(a.message_id).or_default();
|
||||
for d in resolved {
|
||||
if let Some(name) = d["function"]["name"].as_str().map(str::to_string) {
|
||||
if !entry.iter().any(|e| e["function"]["name"].as_str() == Some(name.as_str())) {
|
||||
entry.push(d);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
map
|
||||
}
|
||||
|
||||
/// A **static** catalogue of the MCP servers this user can load — identical
|
||||
/// regardless of which are currently active. Static is deliberate: this text
|
||||
/// sits inside the `cache_control: ephemeral` system block, so the old
|
||||
/// Available/Active split (which moved a server between tables on activation)
|
||||
/// invalidated the prompt-cache prefix on every `activate_tools` call. Which
|
||||
/// servers are active is already visible to the model through the injected
|
||||
/// tools. The grant set is retained in the signature only for call-site
|
||||
/// stability (the DTL serializer may consume it later).
|
||||
fn render_mcp_list(&self, _active_mcp_grants: &HashSet<String>) -> String {
|
||||
let all_servers: std::collections::BTreeSet<String> = self.mcp.tools()
|
||||
.into_iter()
|
||||
.map(|t| t.server_name)
|
||||
@@ -597,37 +683,17 @@ impl MessageBuilder {
|
||||
|
||||
let descriptions = self.mcp.server_descriptions();
|
||||
|
||||
let hidden: Vec<&String> = all_servers.iter()
|
||||
.filter(|n| !active_mcp_grants.contains(*n))
|
||||
.collect();
|
||||
let active: Vec<&String> = all_servers.iter()
|
||||
.filter(|n| active_mcp_grants.contains(*n))
|
||||
.collect();
|
||||
|
||||
let mut out = String::from("## MCP servers\n");
|
||||
|
||||
if !hidden.is_empty() {
|
||||
out.push_str("\n**Available** — call `activate_tools([\"name\"])` to load tools:\n\n");
|
||||
out.push_str("| Server | Description |\n|--------|-------------|\n");
|
||||
for name in &hidden {
|
||||
let desc = descriptions.get(*name)
|
||||
.and_then(|d| d.as_deref())
|
||||
.unwrap_or("—");
|
||||
out.push_str(&format!("| `{name}` | {desc} |\n"));
|
||||
}
|
||||
let mut out = String::from(
|
||||
"## MCP servers\n\nConnectors you can load with `activate_tools([\"name\"])`. \
|
||||
Once loaded, a server's tools are callable as `mcp__<name>__<tool>`:\n\n",
|
||||
);
|
||||
out.push_str("| Server | Description |\n|--------|-------------|\n");
|
||||
for name in &all_servers {
|
||||
let desc = descriptions.get(name)
|
||||
.and_then(|d| d.as_deref())
|
||||
.unwrap_or("—");
|
||||
out.push_str(&format!("| `{name}` | {desc} |\n"));
|
||||
}
|
||||
|
||||
if !active.is_empty() {
|
||||
out.push_str("\n**Active** — tools callable as `mcp__<name>__<tool>`:\n\n");
|
||||
out.push_str("| Server | Description |\n|--------|-------------|\n");
|
||||
for name in &active {
|
||||
let desc = descriptions.get(*name)
|
||||
.and_then(|d| d.as_deref())
|
||||
.unwrap_or("—");
|
||||
out.push_str(&format!("| `{name}` | {desc} |\n"));
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::llm::DtlMode;
|
||||
use super::ChatSessionHandler;
|
||||
use super::message_builder::MessageBuilder;
|
||||
|
||||
@@ -23,6 +24,9 @@ impl ChatSessionHandler {
|
||||
system_substitutions: &HashMap<String, String>,
|
||||
cache_hints: bool,
|
||||
capabilities: &[String],
|
||||
dtl: DtlMode,
|
||||
config_tool_defs: &[Value],
|
||||
activation_stack: Option<i64>,
|
||||
) -> anyhow::Result<Vec<Value>> {
|
||||
let project_root = self.run_context.read().await
|
||||
.as_ref()
|
||||
@@ -45,6 +49,6 @@ impl ChatSessionHandler {
|
||||
// `pool` is passed in from the caller (always `&self.db`) but we take
|
||||
// ownership via Arc::clone above so the signature stays backward-compatible.
|
||||
let _ = pool; // suppress unused-variable warning; MessageBuilder uses its own Arc
|
||||
builder.build(stack_id, agent_id, extra_system_static, extra_system_dynamic, tail_reminder, active_mcp_grants, system_substitutions, cache_hints, capabilities).await
|
||||
builder.build(stack_id, agent_id, extra_system_static, extra_system_dynamic, tail_reminder, active_mcp_grants, system_substitutions, cache_hints, capabilities, dtl, config_tool_defs, activation_stack).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ use std::sync::{Arc, RwLock};
|
||||
|
||||
use anyhow::Result;
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::mcp::McpProvider;
|
||||
use crate::tools::tool_names::CONFIG_GROUP;
|
||||
@@ -16,26 +15,23 @@ use crate::tools::{Tool, ToolDescriptionLength, truncate_label, MAX_LABEL_SHORT}
|
||||
/// - the reserved keyword `"config"` — loads all built-in `Config`-category
|
||||
/// tools (system configuration: MCP/plugin/cron management, secrets).
|
||||
///
|
||||
/// When the LLM calls `activate_tools(["gmail", "config"])`:
|
||||
/// - The in-memory grant set is updated immediately, so the group's tools appear
|
||||
/// in the *next LLM round* of the current turn (via `all_tool_defs()`).
|
||||
/// - If `stack_id` is `None` (root agent): grants are persisted to
|
||||
/// `session_mcp_grants` — they survive across turns and restarts.
|
||||
/// - If `stack_id` is `Some(id)` (sub-agent): grants are persisted to
|
||||
/// `stack_mcp_grants` for that stack frame — they survive restarts but are
|
||||
/// deleted when the frame terminates (`dispatch_call_agent` calls
|
||||
/// `stack_mcp_grants::delete_for_stack` on cleanup).
|
||||
/// When the LLM calls `activate_tools(["gmail", "config"])`, this tool updates
|
||||
/// the in-memory grant set **immediately**, so the group's tools appear in the
|
||||
/// *next LLM round* of the current turn (via `all_tool_defs()`).
|
||||
///
|
||||
/// The `session_mcp_grants` / `stack_mcp_grants` tables store the group string
|
||||
/// verbatim, so `"config"` is persisted just like an MCP server name.
|
||||
/// The **durable** record of the activation is written by the round loop
|
||||
/// (`handle_tool_call`), not here: it anchors the activation to the assistant
|
||||
/// `message_id` that triggered it, in the owner table `activated_tools`
|
||||
/// (`stack_id NULL` = session-scoped root grant; `Some(id)` = sub-agent frame
|
||||
/// grant, deleted on frame exit). Splitting the write this way lets the loop
|
||||
/// supply the `message_id` — the anchor the DTL serializer positions injected
|
||||
/// tool blocks against — which this tool does not have.
|
||||
///
|
||||
/// Not in the global `ToolRegistry` — injected as an `InterfaceTool` in
|
||||
/// `build_agent_config` (root) and `dispatch_call_agent` (sub-agents).
|
||||
/// `build_agent_config` (root) and `build_sub_agent_config` (sub-agents).
|
||||
pub struct ActivateTools {
|
||||
pub pool: Arc<SqlitePool>,
|
||||
pub session_id: i64,
|
||||
/// `None` for root agents (session-scoped grants).
|
||||
/// `Some(stack_id)` for sub-agents (stack-scoped grants, deleted on frame exit).
|
||||
/// `None` for root agents, `Some(stack_id)` for sub-agents. Used only to
|
||||
/// label the confirmation message; the durable scope is decided by the loop.
|
||||
pub stack_id: Option<i64>,
|
||||
pub mcp: Arc<dyn McpProvider>,
|
||||
/// Shared in-memory grant set. Updated in-place on every call so subsequent
|
||||
@@ -97,32 +93,11 @@ impl Tool for ActivateTools {
|
||||
.map(|t| t.server_name.clone())
|
||||
.collect();
|
||||
|
||||
let pool = Arc::clone(&self.pool);
|
||||
let session_id = self.session_id;
|
||||
let stack_id = self.stack_id;
|
||||
let grants_set = Arc::clone(&self.active_mcp_grants);
|
||||
|
||||
// Persist to DB (session-scoped or stack-scoped) and update in-memory set.
|
||||
// The reserved `config` group is stored verbatim, exactly like a server name.
|
||||
tokio::task::block_in_place(|| {
|
||||
tokio::runtime::Handle::current().block_on(async {
|
||||
for name in &names {
|
||||
match stack_id {
|
||||
None => {
|
||||
crate::db::session_mcp_grants::grant(&pool, session_id, name).await?;
|
||||
}
|
||||
Some(sid) => {
|
||||
crate::db::stack_mcp_grants::grant(&pool, sid, name).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::Ok(())
|
||||
})
|
||||
})?;
|
||||
|
||||
// Update in-memory set so the next LLM round sees the new grants.
|
||||
// Update the in-memory set so the next LLM round sees the new grants.
|
||||
// The durable record (with the triggering `message_id`) is written by the
|
||||
// round loop after this tool returns.
|
||||
{
|
||||
let mut set = grants_set.write()
|
||||
let mut set = self.active_mcp_grants.write()
|
||||
.map_err(|_| anyhow::anyhow!("activate_tools: lock poisoned"))?;
|
||||
for name in &names {
|
||||
set.insert(name.clone());
|
||||
@@ -142,7 +117,7 @@ impl Tool for ActivateTools {
|
||||
})
|
||||
.collect();
|
||||
|
||||
let scope = match stack_id {
|
||||
let scope = match self.stack_id {
|
||||
None => "session".to_string(),
|
||||
Some(s) => format!("stack {s}"),
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user