Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion console/src/functions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ fn register_status(iii: &Arc<IIIClient>, config: &Arc<ConsoleConfig>, engine_url
})
.description(
"Return the console worker's runtime knobs: http_port, engine_url, and version.",
),
)
// console-only plumbing; no other worker (e.g. harness) needs to
// discover or call it.
.metadata(serde_json::json!({ "internal": true })),
);
}
79 changes: 76 additions & 3 deletions console/src/proxy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,13 @@
//! between [`axum::extract::ws::Message`] and
//! [`tokio_tungstenite::tungstenite::Message`].
//!
//! The proxy is intentionally dumb: no buffering, no rewriting, no auth.
//! The engine WebSocket and the iii-browser-sdk client on the page do
//! all the framing — this module just shuttles bytes.
//! The proxy is intentionally dumb — no buffering, no auth — with ONE
//! exception: browser→engine `registerfunction` messages get
//! `metadata.internal = true` stamped on (see
//! [`stamp_internal_registration`]). Everything a console page registers
//! is a live-update delivery target for that page, never a discoverable
//! API, and stamping here (not just in the SPA) means stale/cached
//! bundles can't pollute `engine::functions::list` either.

use std::sync::Arc;

Expand Down Expand Up @@ -62,6 +66,13 @@ async fn handle_ws(client: WebSocket, engine_url: Arc<String>) {
}
};
let is_close = matches!(msg, AxumMessage::Close(_));
let msg = match msg {
AxumMessage::Text(t) => match stamp_internal_registration(&t) {
Some(stamped) => AxumMessage::Text(stamped),
None => AxumMessage::Text(t),
},
other => other,
};
if let Some(out) = axum_to_tungstenite(msg) {
if let Err(e) = engine_tx.send(out).await {
tracing::debug!(error = %e, "browser -> engine: engine send error");
Expand Down Expand Up @@ -104,6 +115,34 @@ async fn handle_ws(client: WebSocket, engine_url: Arc<String>) {
}
}

/// If `text` is a wire `registerfunction` message, return a copy with
/// `metadata.internal = true` merged in; `None` means "forward the
/// original untouched" (not a registration, unparseable, or a metadata
/// shape we don't understand).
pub(crate) fn stamp_internal_registration(text: &str) -> Option<String> {
// Fast path: skip the JSON parse for the overwhelming majority of
// frames (invocations, results, stream sends).
if !text.contains("\"registerfunction\"") {
return None;
}
let mut msg: serde_json::Value = serde_json::from_str(text).ok()?;
if msg.get("type").and_then(|t| t.as_str()) != Some("registerfunction") {
return None;
}
let obj = msg.as_object_mut()?;
match obj.get_mut("metadata") {
Some(serde_json::Value::Object(meta)) => {
meta.insert("internal".into(), serde_json::Value::Bool(true));
}
// Unexpected metadata shape — don't rewrite what we don't understand.
Some(_) => return None,
None => {
obj.insert("metadata".into(), serde_json::json!({ "internal": true }));
}
}
serde_json::to_string(&msg).ok()
}

/// Convert an axum `Message` into a tungstenite `Message`. Returns
/// `None` when the variant has no useful tungstenite equivalent.
pub(crate) fn axum_to_tungstenite(msg: AxumMessage) -> Option<TungMessage> {
Expand Down Expand Up @@ -190,6 +229,40 @@ mod tests {
}
}

#[test]
fn stamp_adds_internal_metadata_when_absent() {
let wire = r#"{"type":"registerfunction","id":"console::harness-watch::r0::console-abc"}"#;
let out = stamp_internal_registration(wire).unwrap();
let v: serde_json::Value = serde_json::from_str(&out).unwrap();
assert_eq!(v["metadata"]["internal"], serde_json::json!(true));
assert_eq!(v["id"], "console::harness-watch::r0::console-abc");
}

#[test]
fn stamp_merges_into_existing_metadata() {
let wire = r#"{"type":"registerfunction","id":"x","metadata":{"tenant":"acme"}}"#;
let out = stamp_internal_registration(wire).unwrap();
let v: serde_json::Value = serde_json::from_str(&out).unwrap();
assert_eq!(v["metadata"]["internal"], serde_json::json!(true));
assert_eq!(v["metadata"]["tenant"], "acme");
}

#[test]
fn stamp_ignores_other_messages_and_bad_input() {
// Different message type — even one that mentions registerfunction in a payload.
assert!(stamp_internal_registration(
r#"{"type":"invokefunction","payload":"\"registerfunction\""}"#
)
.is_none());
// Not JSON.
assert!(stamp_internal_registration("registerfunction{").is_none());
// Metadata of an unexpected shape is left alone.
assert!(stamp_internal_registration(
r#"{"type":"registerfunction","id":"x","metadata":"weird"}"#
)
.is_none());
}

#[test]
fn raw_frame_is_dropped() {
// Raw frames are an internal tungstenite construct; we don't
Expand Down
12 changes: 11 additions & 1 deletion console/web/src/lib/iii-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import {
type IIIConnectionState,
type ISdk,
type RegisterFunctionOptions,
type RegisterTriggerInput,
type RemoteFunctionHandler,
registerWorker,
Expand All @@ -42,6 +43,7 @@ export interface IiiClient {
on<P = unknown>(
functionId: string,
handler: (payload: P) => void | Promise<void>,
options?: RegisterFunctionOptions,
): () => void
/**
* Register an engine trigger bound to a function id. Thin passthrough to
Expand Down Expand Up @@ -141,6 +143,7 @@ function wrapSdk(sdk: ISdk, browserId: string): IiiClient {
function on<P>(
functionId: string,
handler: (payload: P) => void | Promise<void>,
options?: RegisterFunctionOptions,
): () => void {
const id = `${functionId}::${browserId}`
// Wrap to satisfy the SDK's RemoteFunctionHandler signature (returns
Expand All @@ -149,7 +152,14 @@ function wrapSdk(sdk: ISdk, browserId: string): IiiClient {
await handler(data as P)
return null
}
const ref = sdk.registerFunction(id, wrapped)
// Every handler registered here is a browser-local console plumbing fn
// (traces, sessions, worktree events, ...) — none are meant for other
// workers (e.g. harness) to discover, so default them out of
// `engine::functions::list`. Callers can still override via `options`.
const ref = sdk.registerFunction(id, wrapped, {
metadata: { internal: true },
...options,
})
Comment on lines +155 to +162

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Shallow spread drops internal: true when caller provides metadata in options.

{ metadata: { internal: true }, ...options } performs a shallow merge. If a caller passes options containing any metadata property, the entire default { internal: true } is replaced — not deep-merged. A caller adding { metadata: { custom: "x" } } would unintentionally lose the internal flag.

The proxy stamping in console/src/proxy.rs provides defense-in-depth, so the practical risk is low. But if you want the SPA default to be robust on its own, consider deep-merging or documenting that callers must include internal: true when passing metadata.

♻️ Optional: deep-merge metadata to preserve the default
     const ref = sdk.registerFunction(id, wrapped, {
-      metadata: { internal: true },
-      ...options,
+      ...options,
+      metadata: { internal: true, ...options?.metadata },
     })
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Every handler registered here is a browser-local console plumbing fn
// (traces, sessions, worktree events, ...) — none are meant for other
// workers (e.g. harness) to discover, so default them out of
// `engine::functions::list`. Callers can still override via `options`.
const ref = sdk.registerFunction(id, wrapped, {
metadata: { internal: true },
...options,
})
// Every handler registered here is a browser-local console plumbing fn
// (traces, sessions, worktree events, ...) — none are meant for other
// workers (e.g. harness) to discover, so default them out of
// `engine::functions::list`. Callers can still override via `options`.
const ref = sdk.registerFunction(id, wrapped, {
...options,
metadata: { internal: true, ...options?.metadata },
})
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@console/web/src/lib/iii-client.ts` around lines 155 - 162, Preserve the
default internal metadata when callers provide metadata through options. In the
handler registration block using sdk.registerFunction, merge options.metadata
with { internal: true } after spreading options, ensuring caller metadata cannot
overwrite the internal flag while retaining other metadata fields.

let active = true
const unregister = () => {
if (!active) return
Expand Down
3 changes: 2 additions & 1 deletion context-manager/src/configuration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,8 @@ pub fn register_config_trigger(
"Internal: hot-reload context-manager from the authoritative configuration when it \
changes — rebuilds the compaction lease store on a lease_dir change and swaps the \
per-call tuning snapshot otherwise.",
),
)
.metadata(json!({ "internal": true })),
);

iii.register_trigger(RegisterTriggerInput {
Expand Down
51 changes: 34 additions & 17 deletions context-manager/src/functions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,12 +71,15 @@ pub(crate) async fn resolve_model(
}

/// Register one typed handler under `id`, mapping `ContextError` into
/// the bus error shape (`code: message`).
/// the bus error shape (`code: message`). `internal` hides the function
/// from the discoverable `engine::functions::list` (trigger/config plumbing
/// stays callable by id); see harness's iii-permissions.yaml.
fn register<Req, Resp, F, Fut>(
iii: &Arc<IIIClient>,
deps: &Arc<Deps>,
id: &str,
description: &str,
internal: bool,
handler: F,
) where
Req: DeserializeOwned + JsonSchema + Send + 'static,
Expand All @@ -85,32 +88,46 @@ fn register<Req, Resp, F, Fut>(
Fut: Future<Output = Result<Resp, ContextError>> + Send + 'static,
{
let deps = deps.clone();
iii.register_function(
id,
RegisterFunction::new_async(move |req: Req| {
let deps = deps.clone();
let handler = handler.clone();
async move { handler(deps, req).await.map_err(Error::from) }
})
.description(description),
);
let reg = RegisterFunction::new_async(move |req: Req| {
let deps = deps.clone();
let handler = handler.clone();
async move { handler(deps, req).await.map_err(Error::from) }
})
.description(description);
let reg = if internal {
reg.metadata(serde_json::json!({ "internal": true }))
} else {
reg
};
iii.register_function(id, reg);
}

pub fn register_all(iii: &Arc<IIIClient>, deps: &Arc<Deps>) {
register(iii, deps, ASSEMBLE_ID, ASSEMBLE_DESC, |d, r| async move {
assemble::handle(&d, r).await
});
register(iii, deps, COMPACT_ID, COMPACT_DESC, |d, r| async move {
compact::handle(&d, r).await
});
register(iii, deps, PRUNE_ID, PRUNE_DESC, |d, r| async move {
register(
iii,
deps,
ASSEMBLE_ID,
ASSEMBLE_DESC,
true,
|d, r| async move { assemble::handle(&d, r).await },
);
register(
iii,
deps,
COMPACT_ID,
COMPACT_DESC,
true,
|d, r| async move { compact::handle(&d, r).await },
);
register(iii, deps, PRUNE_ID, PRUNE_DESC, false, |d, r| async move {
prune::handle(&d, r).await
});
register(
iii,
deps,
COUNT_TOKENS_ID,
COUNT_TOKENS_DESC,
false,
|d, r| async move { count_tokens::handle(&d, r).await },
);

Expand Down
3 changes: 2 additions & 1 deletion harness/src/configuration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,8 @@ pub fn register_config_trigger(
"Internal: hot-reload harness from the authoritative configuration when it changes — \
re-binds the cron pending-sweep on a sweep_expression change and swaps the per-call \
tuning snapshot otherwise.",
),
)
.metadata(json!({ "internal": true })),
);

iii.register_trigger(RegisterTriggerInput {
Expand Down
55 changes: 42 additions & 13 deletions harness/src/functions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,35 @@ fn register<Req, Resp, F, Fut>(
);
}

/// Like [`register`], but tags the registration `metadata.internal = true` so
/// the default `engine::functions::list` hides it: trusted control-plane /
/// loop plumbing, invoked by id, never meant for agent discovery (mirrors the
/// deny rules in iii-permissions.yaml).
fn register_internal<Req, Resp, F, Fut>(
iii: &Arc<IIIClient>,
deps: &Arc<Deps>,
id: &str,
description: &str,
handler: F,
) where
Req: DeserializeOwned + JsonSchema + Send + 'static,
Resp: Serialize + JsonSchema + Send + 'static,
F: Fn(Arc<Deps>, Req) -> Fut + Send + Sync + Clone + 'static,
Fut: Future<Output = Result<Resp, HarnessError>> + Send + 'static,
{
let deps = deps.clone();
iii.register_function(
id,
RegisterFunction::new_async(move |req: Req| {
let deps = deps.clone();
let handler = handler.clone();
async move { handler(deps, req).await.map_err(Error::from) }
})
.description(description)
.metadata(serde_json::json!({ "internal": true })),
);
}

/// Like [`register`], but the handler also receives the per-invocation
/// `metadata` sidecar (`engine::register_trigger`'s `metadata`). Used by the
/// trigger-bridge target `harness::react`.
Expand Down Expand Up @@ -139,41 +168,41 @@ fn register_with_metadata<Req, Resp, F, Fut>(
}

pub fn register_all(iii: &Arc<IIIClient>, deps: &Arc<Deps>) {
register(iii, deps, SEND_ID, SEND_DESC, |d, r| async move {
register_internal(iii, deps, SEND_ID, SEND_DESC, |d, r| async move {
send::handle(&d, r).await
});
register(iii, deps, SPAWN_ID, SPAWN_DESC, |d, r| async move {
spawn::handle(&d, r).await
});
register(iii, deps, TURN_ID, TURN_DESC, |d, r| async move {
register_internal(iii, deps, TURN_ID, TURN_DESC, |d, r| async move {
turn::handle(&d, r).await
});
register(
register_internal(
iii,
deps,
FUNCTION_TRIGGER_ID,
FUNCTION_TRIGGER_DESC,
|d, r| async move { function_trigger::handle(&d, r).await },
);
register(
register_internal(
iii,
deps,
FUNCTION_RESOLVE_ID,
FUNCTION_RESOLVE_DESC,
|d, r| async move { function_resolve::handle(&d, r).await },
);
register(iii, deps, STOP_ID, STOP_DESC, |d, r| async move {
register_internal(iii, deps, STOP_ID, STOP_DESC, |d, r| async move {
stop::handle(&d, r).await
});
register(iii, deps, STATUS_ID, STATUS_DESC, |d, r| async move {
status::handle(&d, r).await
});

// Trusted control-plane (console) — registered, kept off the agent catalog.
register(iii, deps, UNQUEUE_ID, UNQUEUE_DESC, |d, r| async move {
register_internal(iii, deps, UNQUEUE_ID, UNQUEUE_DESC, |d, r| async move {
send::unqueue(&d, r).await
});
register(
register_internal(
iii,
deps,
EDIT_QUEUED_ID,
Expand All @@ -183,28 +212,28 @@ pub fn register_all(iii: &Arc<IIIClient>, deps: &Arc<Deps>) {

// Internal filesystem grant controls — registered for trusted callers, kept
// off the model-facing catalog.
register(
register_internal(
iii,
deps,
FILESYSTEM_GRANT_ID,
FILESYSTEM_GRANT_DESC,
|d, r| async move { filesystem::grant(&d, r).await },
);
register(
register_internal(
iii,
deps,
FILESYSTEM_GRANTS_ID,
FILESYSTEM_GRANTS_DESC,
|d, r| async move { filesystem::grants(&d, r).await },
);
register(
register_internal(
iii,
deps,
FILESYSTEM_REVOKE_ID,
FILESYSTEM_REVOKE_DESC,
|d, r| async move { filesystem::revoke(&d, r).await },
);
register(
register_internal(
iii,
deps,
FILESYSTEM_INFO_ID,
Expand All @@ -213,7 +242,7 @@ pub fn register_all(iii: &Arc<IIIClient>, deps: &Arc<Deps>) {
);

// Internal cron target — registered, but kept off the public catalog.
register(
register_internal(
iii,
deps,
sweep_pending::SWEEP_PENDING_ID,
Expand All @@ -222,7 +251,7 @@ pub fn register_all(iii: &Arc<IIIClient>, deps: &Arc<Deps>) {
);

// Internal session::deleted cleanup — registered, kept off the catalog.
register(
register_internal(
iii,
deps,
crate::subscriptions::ON_SESSION_DELETED_ID,
Expand Down
Loading
Loading