2190 lines · 83.9 KB · Zig
RawDownload
const std = @import("std");const acp_runner = @import("../core/cli/acp_runner.zig");const config_runtime = @import("../core/config/config_runtime.zig");const io_mod = @import("../core/shared/io.zig");const host_target = @import("../core/hosts/target.zig");const jsonrpc = @import("jsonrpc.zig");const acp_types = @import("types.zig");const sessions = @import("sessions.zig");const session_test_controls = @import("session_test_controls.zig");const prompt_handler = @import("prompt.zig");const prompt_test_controls = @import("prompt_test_controls.zig");const app_lifecycle = @import("../core/app/app_lifecycle.zig");const app_runtime_setup = @import("../core/app/app_runtime_setup.zig");const builtin_skills = @import("../builtins/skills.zig");const builtin_tools = @import("../builtins/tools.zig");const credentials = @import("../core/auth/credentials.zig");const debug_trace = @import("../core/shared/debug_trace.zig");const gateway_provider = @import("../core/gateway/gateway_provider.zig");const hooks = @import("../core/hooks/hooks.zig");const mcp_runtime = @import("../core/mcp/mcp_runtime.zig");const mode_registry = @import("../core/modes/mode_registry.zig");const sandbox = @import("../core/permissions/sandbox.zig");const skill_runtime = @import("../core/skills/skill_runtime.zig");const session_codec = @import("../core/session/session_codec.zig");const session_log = @import("../core/session/session_log.zig");const session_store = @import("../core/session/session_store.zig");const session_runtime = @import("../core/session/session.zig");const worker_runtime = @import("../core/agent/worker_runtime.zig");const background_runtime = @import("../core/background/background_runtime.zig");const terminal_client_runtime = @import("../core/terminal/client.zig");const subagent_tool_host = @import("../core/subagent/tool_host.zig");const subagent_authority = @import("../core/subagent/authority.zig");const types = @import("../core/shared/types.zig");const context_contract = @import("../core/workspace/context_contract.zig");const workspace_access = @import("../core/workspace/workspace_access.zig");const web_fetch_runtime = @import("../core/tooling/web_fetch_runtime.zig");const web_search_runtime = @import("../core/tooling/web_search_runtime.zig");const elicitation = @import("../core/mcp/elicitation.zig");const tool_mcp_runtime = @import("../core/tooling/tool_mcp_runtime.zig");const permissions = @import("../core/permissions/permissions.zig"); const Allocator = std.mem.Allocator;const ErrorCode = jsonrpc.ErrorCode;const writeJsonStr = jsonrpc.writeJsonStr;const legacy_url_completion_timeout_ms: i64 = 10 * 60 * 1000; const AcpMethod = enum {    initialize,    session_cancel,    session_new,    session_load,    session_resume,    session_close,    session_list,    session_remove,    session_prompt,    session_set_config_option,    session_set_mode,    unknown,     fn parse(method: []const u8) AcpMethod {        if (std.mem.eql(u8, method, "initialize")) return .initialize;        if (std.mem.eql(u8, method, "session/cancel")) return .session_cancel;        if (std.mem.eql(u8, method, "session/new")) return .session_new;        if (std.mem.eql(u8, method, "session/load")) return .session_load;        if (std.mem.eql(u8, method, "session/resume")) return .session_resume;        if (std.mem.eql(u8, method, "session/close")) return .session_close;        if (std.mem.eql(u8, method, "session/list")) return .session_list;        if (std.mem.eql(u8, method, "session/remove")) return .session_remove;        if (std.mem.eql(u8, method, "session/prompt")) return .session_prompt;        if (std.mem.eql(u8, method, "session/set_config_option")) return .session_set_config_option;        if (std.mem.eql(u8, method, "session/set_mode")) return .session_set_mode;        return .unknown;    }     fn waitsForActivePrompt(self: AcpMethod) bool {        return switch (self) {            .initialize,            .session_cancel,            .session_set_mode,            .session_new,            .session_load,            .session_resume,            .session_close,            => false,            .session_list,            .session_remove,            .session_prompt,            .session_set_config_option,            .unknown,            => true,        };    }}; pub const Config = acp_runner.Config; pub const OutboundKind = enum {    permission,    elicitation,}; pub const OutboundResponse = struct {    result_json: ?[]u8 = null,    error_json: ?[]u8 = null,    cancelled: bool = false,     pub fn deinit(self: *OutboundResponse, alloc: Allocator) void {        if (self.result_json) |value| alloc.free(value);        if (self.error_json) |value| alloc.free(value);        self.* = undefined;    }}; const PendingOutbound = struct {    kind: OutboundKind,    response: ?OutboundResponse = null,}; const max_pending_outbound = 32; const PendingLegacyUrl = struct {    server_name: []u8,    source_id: []u8,    acp_id: []u8,    session_id: []u8,    tool_call_id: []u8,    binding: elicitation.Binding,    accepted: bool = false,    completed: bool = false,     fn deinit(self: *PendingLegacyUrl, alloc: Allocator) void {        alloc.free(self.server_name);        alloc.free(self.source_id);        alloc.free(self.acp_id);        alloc.free(self.session_id);        alloc.free(self.tool_call_id);        self.* = undefined;    }}; pub const ActiveSessionState = struct {    session_id: []u8,    store: ?session_store.Store = null,    writable: ?session_store.LoadedWritableSession = null,    wasm_state: ?session_codec.DurableSessionState = null,    wasm_revision: ?[]u8 = null,    session_write_mutex: std.Io.Mutex = .init,    model: []u8,    mode: []const u8,    workspace_root: []const u8,    api_key: []const u8,    credential_source: ?types.CredentialSource = null,    agent_step_limit: usize,    max_tool_result_bytes: usize,    fast_mode: bool,    effort: types.ReasoningEffort,    first_call_tool_choice: types.ToolChoice,    permission_mode: types.PermissionMode,    sandbox_backend: sandbox.BackendKind,    permission_rules: types.PermissionRuleSet,    /// Runtime-only "allow for this session" grants. Never persisted to    /// profile or project configuration.    session_grants: []types.PermissionGrant = &.{},    session_rt: session_runtime.SessionRuntime,    mcp: ?*mcp_runtime.McpRuntime = null,    cancel_flag: std.atomic.Value(bool),    pending_prompt_id: ?jsonrpc.RequestId,     pub fn retainGrant(self: *ActiveSessionState, alloc: Allocator, tool_name: []const u8, target_path: []const u8) !void {        for (self.session_grants) |grant| {            if (std.mem.eql(u8, grant.tool_name, tool_name) and                std.mem.eql(u8, grant.target_path, target_path)) return;        }         const name_copy = try alloc.dupe(u8, tool_name);        errdefer alloc.free(name_copy);        const target_copy = try alloc.dupe(u8, target_path);        errdefer alloc.free(target_copy);        const next = try alloc.alloc(types.PermissionGrant, self.session_grants.len + 1);        errdefer alloc.free(next);        if (self.session_grants.len > 0) {            std.mem.copyForwards(types.PermissionGrant, next[0..self.session_grants.len], self.session_grants);            alloc.free(self.session_grants);        }        next[next.len - 1] = .{ .tool_name = name_copy, .target_path = target_copy };        self.session_grants = next;    }}; const ActivePrompt = struct {    state: *ServerState,    alloc: Allocator,    msg: jsonrpc.Message,    /// Mode and permission policy captured when the prompt was dispatched.    /// Mid-turn mode changes apply to the next prompt, never the running one.    mode: []const u8,    permission_mode: types.PermissionMode,    sandbox_backend: sandbox.BackendKind,    thread: if (host_target.is_wasm) void else std.Thread = if (host_target.is_wasm) {} else undefined,    reapable: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),}; pub const ServerState = struct {    alloc: Allocator,    cfg: Config,    writer: jsonrpc.Writer,    initialized: bool = false,    client_fs_read: bool = false,    client_fs_write: bool = false,    client_terminal: bool = false,    client_elicitation: elicitation.Capabilities = .{},    workspace_root: []u8 = &.{},    workspace_access: workspace_access.WorkspaceAccess = .{},    api_key: []u8 = &.{},    credential_source: ?types.CredentialSource = null,    gateway_team: ?[]u8 = null,    selected_model: []u8 = &.{},    configured_model: []u8 = &.{},    process_model_override: bool = false,    permission_mode: types.PermissionMode = .ask,    permission_rules: types.PermissionRuleSet = .{},    agent_step_limit: usize = 0,    max_tool_result_bytes: usize = 64 * 1024,    context_limits: config_runtime.context_limits.Values = .{},    fast_mode: bool = false,    effort: types.ReasoningEffort = .auto,    first_call_tool_choice: types.ToolChoice = .auto,    context_enabled: bool = true,    sandbox_backend: sandbox.BackendKind = .none,    active_session: ?ActiveSessionState = null,    active_prompt: ?*ActivePrompt = null,    subagent_authority_mutex: std.Io.Mutex = .init,    skills: skill_runtime.Runtime = .{},    context_snapshot: context_contract.GatheredContextSnapshot = .{},    worker: worker_runtime.WorkerRuntime = .{},    background: background_runtime.BackgroundRuntime = .{},    terminal_client: terminal_client_runtime.Runtime = .{},    subagent_store: ?session_store.Store = null,    subagent_host: ?*subagent_tool_host.Runtime = null,    capability_resolver: gateway_provider.CapabilityResolver = .{},    terminate_connection: bool = false,    web_fetch_runtime: web_fetch_runtime.Runtime = web_fetch_runtime.Runtime.init(.{}),    web_search_runtime: web_search_runtime.Runtime = web_search_runtime.Runtime.init(.{}),    lifecycle_runtime: hooks.Runtime = hooks.Runtime.init(std.heap.c_allocator),    lifecycle_view: hooks.RuntimeView = hooks.RuntimeView.empty(),    outbound_mutex: std.Io.Mutex = .init,    outbound_cond: std.Io.Condition = .init,    next_outbound_request_id: u64 = 1,    pending_outbound: std.AutoHashMapUnmanaged(u64, PendingOutbound) = .empty,    legacy_url_mutex: std.Io.Mutex = .init,    pending_legacy_urls: std.ArrayListUnmanaged(PendingLegacyUrl) = .empty,     pub fn deinit(self: *ServerState) void {        reapActivePrompt(self, true);        self.terminal_client.deinit();        closeActiveSession(self) catch |err| {            debug_trace.logf(                "session",                "failed to flush ACP session usage during shutdown err={s}",                .{@errorName(err)},            );        };        self.workspace_access.deinit(self.alloc);        if (self.workspace_root.len > 0) self.alloc.free(self.workspace_root);        if (self.api_key.len > 0) self.alloc.free(self.api_key);        if (self.gateway_team) |team| self.alloc.free(team);        if (self.selected_model.len > 0) self.alloc.free(self.selected_model);        if (self.configured_model.len > 0) self.alloc.free(self.configured_model);        self.permission_rules.deinit(self.alloc);        self.background.deinit(std.heap.c_allocator);        self.skills.deinit(self.alloc);        self.context_snapshot.deinit(self.alloc);        self.worker.deinit(std.heap.c_allocator);        self.web_fetch_runtime.deinit(self.alloc);        self.web_search_runtime.deinit();        self.lifecycle_runtime.deinit();        self.capability_resolver.deinit(self.alloc);        var pending = self.pending_outbound.valueIterator();        while (pending.next()) |entry| {            if (entry.response) |*response| response.deinit(self.alloc);        }        self.pending_outbound.deinit(self.alloc);        clearPendingLegacyUrls(self);        self.pending_legacy_urls.deinit(self.alloc);    }}; pub fn releaseActiveSession(state: *ServerState) !void {    clearPendingLegacyUrls(state);    const active = if (state.active_session) |*session| session else return;    disableSubagentHost(state);    if (comptime !host_target.is_wasm) {        active.session_rt.usage.cancelReconciliation();        active.session_rt.usage.finishProfilePublicationsBeforeShutdown();        flushActiveSessionUsage(state) catch |err| {            active.session_rt.usage.startReconciliation(                state.alloc,                state.api_key,            );            return err;        };        active.session_rt.usage.configurePublicationSink(null);        active.session_rt.usage.configureCheckpointSink(null);    }    destroyActiveSession(state);} fn closeActiveSession(state: *ServerState) !void {    const active = if (state.active_session) |*session| session else return;    disableSubagentHost(state);    if (comptime !host_target.is_wasm) {        active.session_rt.usage.cancelReconciliation();        active.session_rt.usage.finishProfilePublicationsBeforeShutdown();        flushActiveSessionUsage(state) catch |err| {            destroyActiveSession(state);            return err;        };        active.session_rt.usage.configurePublicationSink(null);        active.session_rt.usage.configureCheckpointSink(null);    }    destroyActiveSession(state);} fn destroyActiveSession(state: *ServerState) void {    const active = if (state.active_session) |*session| session else return;    state.background.detachManagedPersistence(        std.heap.c_allocator,        active.session_id,    );    state.alloc.free(active.session_id);    state.alloc.free(active.model);    types.freePermissionGrantSlice(state.alloc, active.session_grants);    if (comptime !host_target.is_wasm) {        if (active.mcp) |runtime| {            runtime.retireAndWait();            runtime.deinit();            state.alloc.destroy(runtime);        }    }    active.session_rt.deinit(state.alloc);    if (active.writable) |*writable| writable.deinit(state.alloc);    if (active.store) |*store| store.deinit(state.alloc);    if (active.wasm_state) |*wasm_state| wasm_state.deinit(state.alloc);    if (active.wasm_revision) |revision| state.alloc.free(revision);    state.active_session = null;} pub fn enableSubagentHost(state: *ServerState) void {    disableSubagentHost(state);    const active = if (state.active_session) |*session| session else return;    if (active.writable == null) return;    state.subagent_store = session_store.Store.init(state.alloc, state.workspace_root) catch |err| {        debug_trace.logf("acp", "subagent host store unavailable session={s} err={s}", .{ active.session_id, @errorName(err) });        return;    };    state.subagent_host = subagent_tool_host.Runtime.create(        state.alloc,        &state.subagent_store.?,        active.session_id,        .{ .context = state, .resolve_fn = resolveSubagentAuthority },        .{ .context = state, .run_fn = prompt_handler.runSubagentChild },    ) catch |err| {        debug_trace.logf("acp", "subagent host unavailable session={s} err={s}", .{ active.session_id, @errorName(err) });        state.subagent_store.?.deinit(state.alloc);        state.subagent_store = null;        return;    };    state.subagent_host.?.requestBackgroundRecovery(        io_mod.milliTimestamp(),    ) catch |err| {        debug_trace.logf(            "subagent",            "acp background recovery unavailable root_id={s} outcome={s}",            .{ state.subagent_host.?.root_id, @errorName(err) },        );    };} fn resolveSubagentAuthority(    raw: ?*anyopaque,    alloc: Allocator,    root_id: []const u8,) subagent_authority.HostResolveError!subagent_authority.HostAuthority {    const state: *ServerState = @ptrCast(@alignCast(raw.?));    const active = if (state.active_session) |*value| value else return error.HostAuthorityUnavailable;    if (!std.mem.eql(u8, active.session_id, root_id)) {        return error.HostAuthorityUnavailable;    }    state.subagent_authority_mutex.lockUncancelable(io_mod.getIo());    defer state.subagent_authority_mutex.unlock(io_mod.getIo());    const integrations = if (active.mcp) |mcp|        mcp.snapshotToolNames(alloc, active.permission_rules)    else        alloc.alloc([]u8, 0);    const owned_integrations = integrations catch return error.OutOfMemory;    defer {        for (owned_integrations) |name| alloc.free(name);        alloc.free(owned_integrations);    }    var mcp_view = if (active.mcp) |mcp|        try mcp.snapshotAccessView(            alloc,            root_id,            root_id,            active.permission_rules,            state.cfg.mode_registry.toolAllowed(builtin_tools.advertisement_set, active.mode, "mcp_features") and                !permissions.rulesDenyAllTargetsForTool(active.permission_rules, "mcp_features"),        )    else        null;    defer if (mcp_view) |*view| view.deinit(alloc);    var permission_state = active.session_rt.snapshotPermissionState(alloc) catch |err| switch (err) {        error.OutOfMemory => return error.OutOfMemory,        else => return error.HostAuthorityUnavailable,    };    defer permission_state.deinit(alloc);    return subagent_tool_host.captureHostAuthorityWithMcpView(        alloc,        .{            .tool_set = builtin_tools.advertisement_set,            .mode = .{                .active = .{                    .registry = state.cfg.mode_registry,                    .id = active.mode,                },            },        },        active.sandbox_backend,        owned_integrations,        active.permission_rules,        active.session_grants,        permission_state,        if (mcp_view) |*view| view else null,    );} pub fn disableSubagentHost(state: *ServerState) void {    if (state.subagent_host) |host| {        host.deinit();        state.subagent_host = null;    }    if (state.subagent_store) |*store| {        store.deinit(state.alloc);        state.subagent_store = null;    }} fn flushActiveSessionUsage(state: *ServerState) !void {    const active = if (state.active_session) |*session| session else return;    const writable = if (active.writable) |*value| value else return;    if (!writable.needsFinalStateReplacement(        active.session_rt.usage.isDirty(),    )) return;     var current = try writable.state.dupe(state.alloc);    defer current.deinit(state.alloc);    const history = try active.session_rt.snapshotHistory(state.alloc);    types.freeHistoryTurnSlice(state.alloc, current.history);    current.history = history;    const permission_state = try active.session_rt.snapshotPermissionState(state.alloc);    current.permission_state.deinit(state.alloc);    current.permission_state = permission_state;    current.conversation_language = active.session_rt.languageSnapshot();    const usage_snapshot = try active.session_rt.usage.snapshot(state.alloc);    if (current.usage) |*old| old.deinit(state.alloc);    current.usage = usage_snapshot;    const store = if (active.store) |*value|        value    else        return error.SessionPersistenceUnavailable;    const recovery_checkpoint = try store.prepareUsageRecoveryCheckpoint(        state.alloc,        writable,        usage_snapshot,    );    current.updated_at_ms = recovery_checkpoint.timestamp_ms;    _ = try writable.commitStateReplacement(        state.alloc,        current,        .compaction,        .retry_expected_tail,        .{},    );    try store.finishUsageRecoveryCheckpoint(        writable.active_id,        recovery_checkpoint,    );    if (current.usage) |usage| {        active.session_rt.usage.markClean(usage);    }} pub fn run(alloc: Allocator, cfg: Config) !void {    return runWithTransport(alloc, cfg, jsonrpc.Reader.init(), jsonrpc.Writer.init());} pub fn runWithTransport(    alloc: Allocator,    cfg: Config,

Lines 1–500 of 2190.