@aexhq/sdk 0.58.0 → 0.59.1

package/dist/session.js

@@ -1,15 +1,15 @@
-import { CustomerEnvironment } from "@aexhq/session-protocol";
-import { inspectEnvironment } from "@aexhq/environment";
+import { CustomerEnvironment, prepareComponents, } from "@aexhq/brain";
 import * as z from "zod";
 import { AbortError, OutputRefusalError, OutputSchemaError, OutputValidationError, SessionError, abortError, errorFromApi, } from "./errors.js";
 import { canonicalize, jcsSha256, randomIdempotencyKey } from "./json.js";
-import { compileTools } from "./tools.js";
-import { encodeBase64, SessionChildren, SessionStorage } from "./resources.js";
+import { compileCallbacks } from "./tools.js";
+import { SessionChildren, SessionSandbox, SessionStorage } from "./resources.js";
 export class Sessions {
     #transport;
     #webSocketFactory;
     #customerEnvironments = new Map();
     #customerEnvironmentInstances = new Map();
+    #customerEnvironmentReadinessWaiters = new Map();
     #closed = false;
     constructor(transport, webSocketFactory) {
         this.#transport = transport;
@@ -20,21 +20,109 @@ export class Sessions {
         if (this.#closed)
             return;
         this.#closed = true;
-        for (const environment of this.#customerEnvironmentInstances.values())
-            environment.close();
+        for (const hand of this.#customerEnvironmentInstances.values())
+            hand.close();
         this.#customerEnvironmentInstances.clear();
         this.#customerEnvironments.clear();
     }
     async create(options, request = {}) {
         if (this.#closed)
             throw new SessionError("Aex client is closed");
-        if (options.loop === undefined)
-            throw new TypeError("sessions.create requires an imported loop");
-        const compiledTools = await compileTools(options.tools, options.environments, options.secrets);
-        await Promise.all(compiledTools.layers.map((layer) => this.#transport.ensureArtifactLayer(layer.checksum, layer.media_type, layer.content, request.signal)));
-        await this.#ensureCustomerEnvironment(compiledTools.callbackClientId, compiledTools.clientRegistrations, request.signal);
+        if (options.model?.component === undefined) {
+            throw new TypeError("sessions.create requires an imported Model component");
+        }
+        if (options.agentloop === undefined) {
+            throw new TypeError("sessions.create requires an imported Agentloop component");
+        }
+        const model = options.model.component;
+        const agentloop = options.agentloop;
+        const tools = [...(options.tools ?? [])];
+        const componentTools = tools.filter(isComponentTool);
+        const callbackTools = tools.filter((tool) => !isComponentTool(tool));
+        const environments = Object.entries(options.environments ?? {});
+        if (environments.some(([, environment]) => !isEnvironmentComponent(environment))) {
+            throw new TypeError("component sessions accept only precompiled Environment component values");
+        }
+        const environmentComponents = environments.map(([name, environment]) => {
+            if (!isEnvironmentComponent(environment)) {
+                throw new TypeError(`Environment ${JSON.stringify(name)} is not a component`);
+            }
+            return [name, environment];
+        });
+        const applicationEnvironments = environmentComponents.filter(([, environment]) => isApplicationEnvironment(environment));
+        if (callbackTools.length > 0 && applicationEnvironments.length !== 1) {
+            throw new TypeError("application callback Tools require exactly one app() Environment");
+        }
+        const callbacks = await compileCallbacks(callbackTools);
+        const allComponentTools = [...componentTools, ...callbacks.components];
+        const toolDefinitions = allComponentTools.map(componentToolDefinition);
+        const names = new Set();
+        for (const definition of toolDefinitions) {
+            if (names.has(definition.name)) {
+                throw new TypeError(`Tool ${JSON.stringify(definition.name)} was selected twice`);
+            }
+            names.add(definition.name);
+        }
+        const callbackEnvironment = applicationEnvironments[0];
+        if (callbackEnvironment !== undefined) {
+            await this.#ensureCustomerEnvironment(applicationEnvironmentId(callbackEnvironment[1]), callbacks.registrations, request.signal);
+        }
+        const prepared = await prepareComponents([
+            model,
+            agentloop,
+            ...allComponentTools,
+            ...environmentComponents.map(([, environment]) => environment),
+        ]);
+        const modelBinding = prepared.bindings[0];
+        const agentloopBinding = prepared.bindings[1];
+        if (modelBinding === undefined || agentloopBinding === undefined) {
+            throw new TypeError("component session bindings are incomplete");
+        }
+        const toolBindings = prepared.bindings.slice(2, 2 + allComponentTools.length);
+        const environmentBindings = prepared.bindings.slice(2 + allComponentTools.length);
+        const toolItems = allComponentTools.map((tool, index) => {
+            const binding = toolBindings[index];
+            if (binding === undefined)
+                throw new TypeError("Tool component binding is missing");
+            const definition = toolDefinitions[index];
+            if (definition === undefined)
+                throw new TypeError("Tool component definition is missing");
+            const needsEnvironment = binding.grants.includes("environment");
+            const isCallback = index >= componentTools.length;
+            if (needsEnvironment && !isCallback && environmentComponents.length !== 1) {
+                throw new TypeError("a Tool with the environment grant requires exactly one declared Environment");
+            }
+            const environmentName = isCallback
+                ? callbackEnvironment?.[0]
+                : environmentComponents[0]?.[0];
+            return {
+                definition,
+                executor: {
+                    kind: "component",
+                    component_digest: binding.component_digest,
+                    world: binding.world,
+                    config: binding.config,
+                    grants: binding.grants,
+                    ...(needsEnvironment ? { environment: environmentName } : {}),
+                },
+            };
+        });
+        const environmentConfig = Object.fromEntries(environmentComponents.map(([name], index) => {
+            const binding = environmentBindings[index];
+            if (binding === undefined) {
+                throw new TypeError(`Environment component binding ${JSON.stringify(name)} is missing`);
+            }
+            return [name, {
+                    component_digest: binding.component_digest,
+                    world: binding.world,
+                    config: binding.config,
+                }];
+        }));
         const body = {
             model: {
+                component_digest: modelBinding.component_digest,
+                world: modelBinding.world,
+                config: modelBinding.config,
                 provider: options.model.provider,
                 name: options.model.name,
                 api_key: options.model.apiKey,
@@ -42,37 +130,24 @@ export class Sessions {
                 ...(options.model.maxOutputTokens === undefined
                     ? {}
                     : { max_output_tokens: options.model.maxOutputTokens }),
-                ...(options.model.contextWindowTokens === undefined
-                    ? {}
-                    : { context_window_tokens: options.model.contextWindowTokens }),
             },
-            tools: {
-                items: compiledTools.items,
-            },
-            ...(Object.keys(compiledTools.environments).length === 0
-                ? {}
-                : { environments: compiledTools.environments }),
-            ...(compiledTools.bundles.length === 0 ? {} : { tool_bundles: compiledTools.bundles }),
             agentloop: {
-                source_bundle_sha256: options.loop.sha256,
-                toolchain: options.loop.toolchain,
-                bundle_base64: encodeBase64(new TextEncoder().encode(options.loop.source)),
+                component_digest: agentloopBinding.component_digest,
+                world: agentloopBinding.world,
+                config: agentloopBinding.config,
             },
+            component_artifacts: prepared.artifacts,
+            tools: { items: toolItems },
+            ...(environmentComponents.length === 0 ? {} : { environments: environmentConfig }),
             ...(options.secrets === undefined ? {} : { secrets: options.secrets }),
             ...(options.metadata === undefined ? {} : { metadata: options.metadata }),
-            ...(options.network === undefined
+            ...(options.retainUntil === undefined
                 ? {}
-                : { network: options.network }),
+                : { retain_until: normalizeTimestamp(options.retainUntil, "retainUntil") }),
+            ...(options.network === undefined ? {} : { network: options.network }),
             ...(options.providerRecoveryRetries === undefined
                 ? {}
                 : { provider_recovery_retries: options.providerRecoveryRetries }),
-            ...(compiledTools.callbackClientId === undefined
-                ? {}
-                : {
-                    client: {
-                        id: compiledTools.callbackClientId,
-                    },
-                }),
             ...(options.children === undefined
                 ? {}
                 : {
@@ -95,7 +170,7 @@ export class Sessions {
             signal: request.signal,
             retry: true,
         });
-        return new Session(this.#transport, data, compiledTools.environmentNames);
+        return new Session(this.#transport, data);
     }
     async get(id, options = {}) {
         const data = await this.#transport.json("GET", `/v1/sessions/${encodeURIComponent(id)}`, { signal: options.signal });
@@ -163,16 +238,78 @@ export class Sessions {
                     }
                 }
             });
-            // The request may stop waiting, but the process-scoped runner remains reconnectable for
-            // later sessions. Its grant/reconnect lifetime must never inherit one create signal.
-            await waitWithSignal(starting, signal);
+            await this.#waitForCustomerEnvironmentReadiness(clientId, starting, signal);
             return;
         }
-        const environment = await waitWithSignal(existing, signal);
+        const hand = await this.#waitForCustomerEnvironmentReadiness(clientId, existing, signal);
         if (this.#closed)
             throw new SessionError("Aex client is closed");
-        await waitWithSignal(environment.register(registrations), signal);
+        await waitWithSignal(hand.register(registrations), signal);
     }
+    async #waitForCustomerEnvironmentReadiness(clientId, starting, signal) {
+        this.#customerEnvironmentReadinessWaiters.set(clientId, (this.#customerEnvironmentReadinessWaiters.get(clientId) ?? 0) + 1);
+        let settled = false;
+        void starting.then(() => { settled = true; }, () => { settled = true; });
+        try {
+            return await waitWithSignal(starting, signal);
+        }
+        catch (error) {
+            if (signal?.aborted && !settled
+                && this.#customerEnvironmentReadinessWaiters.get(clientId) === 1
+                && this.#customerEnvironments.get(clientId) === starting) {
+                const partial = this.#customerEnvironmentInstances.get(clientId);
+                partial?.close();
+                this.#customerEnvironments.delete(clientId);
+                if (this.#customerEnvironmentInstances.get(clientId) === partial) {
+                    this.#customerEnvironmentInstances.delete(clientId);
+                }
+            }
+            throw error;
+        }
+        finally {
+            const waiters = this.#customerEnvironmentReadinessWaiters.get(clientId) ?? 1;
+            if (waiters <= 1)
+                this.#customerEnvironmentReadinessWaiters.delete(clientId);
+            else
+                this.#customerEnvironmentReadinessWaiters.set(clientId, waiters - 1);
+        }
+    }
+}
+function isApplicationEnvironment(value) {
+    const config = value.config;
+    return config !== null && typeof config === "object" && !Array.isArray(config) &&
+        config.driver === "customer";
+}
+function applicationEnvironmentId(value) {
+    const config = value.config;
+    const id = config.configuration?.registration;
+    if (typeof id !== "string" || !/^[A-Za-z0-9_.:-]{1,128}$/u.test(id)) {
+        throw new TypeError("app() Environment registration is invalid");
+    }
+    return id;
+}
+function isComponentTool(value) {
+    return value !== null
+        && typeof value === "object"
+        && value.kind === "brain.component"
+        && value.extension === "tool";
+}
+function isEnvironmentComponent(value) {
+    return value !== null
+        && typeof value === "object"
+        && value.kind === "brain.component"
+        && value.extension === "environment";
+}
+function componentToolDefinition(value) {
+    const config = value.config;
+    if (config === null || typeof config !== "object" || Array.isArray(config)) {
+        throw new TypeError("Tool component config must be an object containing definition");
+    }
+    const definition = config.definition;
+    if (definition === null || typeof definition !== "object" || Array.isArray(definition)) {
+        throw new TypeError("Tool component config.definition is required");
+    }
+    return definition;
 }
 function waitWithSignal(promise, signal) {
     if (signal === undefined)
@@ -192,32 +329,16 @@ function waitWithSignal(promise, signal) {
 export class Session {
     #transport;
     #data;
+    sandbox;
     storage;
     children;
-    #environmentNames;
-    #environmentHandles = new Map();
-    constructor(transport, data, environmentNames = new Map()) {
+    constructor(transport, data) {
         this.#transport = transport;
         this.#data = data;
-        this.#environmentNames = environmentNames;
-        this.storage = new SessionStorage(transport, data.id, environmentNames);
+        this.sandbox = new SessionSandbox(transport, data.id);
+        this.storage = new SessionStorage(transport, data.id);
         this.children = new SessionChildren(transport, data.id);
     }
-    environment(environment) {
-        const name = this.#environmentNames.get(environment);
-        if (name === undefined)
-            throw new TypeError("EnvironmentRef does not belong to this Session");
-        const existing = this.#environmentHandles.get(environment);
-        if (existing !== undefined)
-            return existing;
-        const handle = inspectEnvironment(environment).createHandle({
-            sessionId: this.id,
-            environment: name,
-            request: async (method, path, body) => await this.#transport.json(method, path, body === undefined ? {} : { body }),
-        });
-        this.#environmentHandles.set(environment, handle);
-        return handle;
-    }
     get id() {
         return this.#data.id;
     }
@@ -250,6 +371,9 @@ export class Session {
     get updatedAt() {
         return this.#data.updated_at;
     }
+    get retainUntil() {
+        return this.#data.retain_until;
+    }
     get metadata() {
         return this.#data.metadata;
     }
@@ -351,6 +475,25 @@ export class Session {
         this.#data = await this.#transport.json("POST", `/v1/sessions/${encodeURIComponent(this.id)}/cancel`, { signal: options.signal, retry: true });
         return this;
     }
+    async suspend(options = {}) {
+        this.#data = await this.#transport.json("POST", `/v1/sessions/${encodeURIComponent(this.id)}/suspend`, { signal: options.signal, retry: true });
+        return this;
+    }
+    async resume(options = {}) {
+        this.#data = await this.#transport.json("POST", `/v1/sessions/${encodeURIComponent(this.id)}/resume`, { signal: options.signal, retry: true });
+        return this;
+    }
+    async setRetention(value, options = {}) {
+        this.#data = await this.#transport.json("POST", `/v1/sessions/${encodeURIComponent(this.id)}/retention`, {
+            body: {
+                retain_until: normalizeTimestamp(value, "retainUntil"),
+                allow_shorten: options.allowShorten ?? false,
+            },
+            signal: options.signal,
+            retry: true,
+        });
+        return this;
+    }
     async end(options = {}) {
         this.#data = await this.#transport.json("POST", `/v1/sessions/${encodeURIComponent(this.id)}/end`, { signal: options.signal, retry: true });
         return this;
@@ -369,6 +512,12 @@ export class Session {
 function isOutputOptions(options) {
     return "output" in options;
 }
+function normalizeTimestamp(value, field) {
+    const parsed = value instanceof Date ? value : new Date(value);
+    if (!Number.isFinite(parsed.getTime()))
+        throw new TypeError(`${field} must be a valid timestamp`);
+    return parsed.toISOString();
+}
 async function compileOutputSchema(schema, retries) {
     let jsonSchema;
     try {

package/dist/tools.d.ts

@@ -1,4 +1,5 @@
-import { type ComputerProfile, type EnvironmentRef } from "@aexhq/environment";
+import type { ComponentExtension } from "@aexhq/brain";
+import type { ToolDefinition } from "@aexhq/brain/session";
 import * as z from "zod";
 export type JsonSchema = Record<string, unknown>;
 export interface ToolContext {
@@ -6,37 +7,8 @@ export interface ToolContext {
     readonly operationId: string;
     readonly sessionId: string;
     readonly deadlineMs: number;
-    readonly workspace?: string;
 }
 export type ToolHandler<Input extends z.ZodType, Output = unknown> = (input: z.output<Input>, context: ToolContext) => Output | Promise<Output>;
-export type ToolSetupHandler = () => void | Promise<void>;
-export interface ToolRequirements {
-    readonly env?: readonly string[];
-    readonly workspace?: boolean;
-    readonly processes?: boolean;
-    readonly network?: readonly {
-        readonly host: string;
-        readonly port: number;
-    }[];
-    readonly streaming?: boolean;
-    readonly recovery?: "retained" | "connection" | "replay-safe";
-}
-export interface PreparedArtifact {
-    readonly digest: string;
-    readonly target: ComputerProfile["platform"];
-    readonly bytes: number;
-    readonly execute: string;
-    readonly setup?: string;
-    readonly layers: readonly PreparedArtifactLayer[];
-}
-export interface PreparedArtifactLayer {
-    readonly digest: string;
-    readonly bytes: number;
-    readonly mediaType: "application/javascript+esm" | "application/x-xz";
-    readonly mountPath: string;
-    readonly unpack: "file" | "tar.xz";
-    readonly source: URL;
-}
 export interface ToolContract {
     readonly name: string;
     readonly description?: string;
@@ -50,29 +22,14 @@ export interface Tool<Input extends z.ZodType = z.ZodType, Output = unknown> {
     readonly description?: string;
     readonly input: Input;
     readonly output?: z.ZodType;
-    readonly requirements: ToolRequirements;
     readonly handler: ToolHandler<Input, Output>;
-    readonly setupHandler?: ToolSetupHandler;
-    readonly artifact?: PreparedArtifact;
     describe(description: string): Tool<Input, Output>;
     named(name: string): Tool<Input, Output>;
     returns<Schema extends z.ZodType>(schema: Schema): Tool<Input, z.output<Schema>>;
-    needs(requirements: ToolRequirements): Tool<Input, Output>;
-    setup(handler: ToolSetupHandler): Tool<Input, Output>;
-    bind<Environment extends EnvironmentRef>(environment: Environment): BoundTool<Environment, Input, Output>;
-}
-export interface BoundTool<Environment extends EnvironmentRef = EnvironmentRef, Input extends z.ZodType = z.ZodType, Output = unknown> {
-    readonly kind: "aex.bound-tool";
-    readonly tool: Tool<Input, Output>;
-    readonly environment: Environment;
 }
 declare const EMPTY_INPUT: z.ZodObject<{}, z.core.$strip>;
 export declare function tool<Output>(handler: ToolHandler<typeof EMPTY_INPUT, Output>): Tool<typeof EMPTY_INPUT, Output>;
 export declare function tool<Input extends z.ZodType, Output>(input: Input, handler: ToolHandler<Input, Output>): Tool<Input, Output>;
-export declare function withPreparedArtifact<Input extends z.ZodType, Output>(value: Tool<Input, Output>, artifact: PreparedArtifact): Tool<Input, Output>;
-export type EnvironmentMap = Record<string, EnvironmentRef>;
-export type EnvironmentValue<Environments extends EnvironmentMap> = Environments[keyof Environments];
-export type ToolSelection<Environment extends EnvironmentRef = EnvironmentRef> = Tool | BoundTool<Environment>;
 export interface ClientRegistration {
     readonly registration: string;
     readonly name: string;
@@ -81,20 +38,11 @@ export interface ClientRegistration {
     readonly output?: z.ZodType;
     readonly handler: ToolHandler<z.ZodType>;
 }
-export interface CompiledTools {
-    readonly items: readonly unknown[];
-    readonly bundles: readonly unknown[];
-    readonly layers: readonly CompiledArtifactLayer[];
-    readonly clientRegistrations: readonly ClientRegistration[];
-    readonly environments: Readonly<Record<string, unknown>>;
-    readonly environmentNames: ReadonlyMap<EnvironmentRef, string>;
-    readonly callbackClientId?: string;
-}
-export interface CompiledArtifactLayer {
-    readonly checksum: string;
-    readonly bytes: number;
-    readonly media_type: string;
-    readonly content: Uint8Array;
+export interface CompiledCallbacks {
+    readonly components: readonly ComponentExtension<"tool", Readonly<Record<string, unknown>> & {
+        readonly definition: ToolDefinition;
+    }>[];
+    readonly registrations: readonly ClientRegistration[];
 }
-export declare function compileTools(selections: readonly ToolSelection[] | undefined, declared: EnvironmentMap | undefined, secrets?: Readonly<Record<string, string>>): Promise<CompiledTools>;
+export declare function compileCallbacks(selections: readonly Tool[]): Promise<CompiledCallbacks>;
 export {};