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:
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user