@gea-ai/agent-sdk 0.1.260922-alpha.0 → 0.1.260922-alpha.2

package/dist/agent-session-execution.js

@@ -1,6 +1,6 @@
 import { getSessionPayload, putSessionPayload } from "./agent-session-payload.js";
 import { agentModelCallsSchema, } from "@gea-ai/contract/agent-usage";
-import { AGENT_SESSION_COMPUTER_LEASE_MS, MAX_AGENT_SESSION_TRANSFER_BYTES, agentSessionTransferChunkSchema, agentSessionTransferReferenceSchema, agentSessionTransferIdentitySchema, agentSessionComputerSchema, agentSessionComputerReferenceSchema, AGENT_SESSION_EXECUTION_LEASE_MS, AGENT_SESSION_EXECUTION_PROTOCOL_VERSION, MAX_AGENT_SESSION_RESPONSE_BYTES, MAX_AGENT_SESSION_EXECUTION_ATTEMPTS, MAX_AGENT_SESSION_EXECUTION_BODY_BYTES, MAX_AGENT_SESSION_PAGE_SIZE, agentSessionAttemptIdentitySchema, agentSessionExecutionRenewSchema, agentSessionExecutionCompleteSchema, agentSessionExecutionPauseSchema, agentSessionExecutionControlSchema, agentSessionExecutionMutationSchema, agentSessionExecutionStartSchema, } from "@gea-ai/contract/agent-session";
+import { AGENT_SESSION_COMPUTER_LEASE_MS, MAX_AGENT_SESSION_TRANSFER_BYTES, agentSessionTransferChunkSchema, agentSessionTransferReferenceSchema, agentSessionTransferIdentitySchema, agentSessionComputerSchema, agentSessionComputerReferenceSchema, AGENT_SESSION_EXECUTION_LEASE_MS, AGENT_SESSION_EXECUTION_PROTOCOL_VERSION, MAX_AGENT_SESSION_RESPONSE_BYTES, MAX_AGENT_SESSION_EXECUTION_ATTEMPTS, MAX_AGENT_SESSION_EXECUTION_BODY_BYTES, MAX_AGENT_SESSION_PAGE_SIZE, agentSessionAttemptIdentitySchema, agentSessionExecutionRenewSchema, agentSessionExecutionCompleteSchema, agentSessionExecutionPauseSchema, agentSessionExecutionControlSchema, agentSessionExecutionMutationSchema, agentSessionExecutionStartSchema, agentSessionExecutionPrepareSchema, } from "@gea-ai/contract/agent-session";
 import { z } from "zod";
 import { activeComputerOperation } from "./agent-session-computer.js";
 import { isAgentSessionInboxBusy, scheduleAgentSessionInbox, getAgentSessionRunState, handleAgentSessionInboxRequest, hasAgentSessionInbox, skipAgentSessionInputWait, } from "./agent-session-inbox.js";
@@ -335,17 +335,28 @@ export async function handleAgentSessionExecutionRequest(request, storage) {
                 return { computer };
             }));
         }
+        const capabilities = {
+            protocolVersion: AGENT_SESSION_EXECUTION_PROTOCOL_VERSION,
+            managedExecutionVersion: 1,
+            logicalRunVersion: 1,
+        };
         if (request.method === "GET" && url.pathname === "/v2/capabilities")
-            return json({
-                protocolVersion: AGENT_SESSION_EXECUTION_PROTOCOL_VERSION,
-                managedExecutionVersion: 1,
-                logicalRunVersion: 1,
-            });
+            return json(capabilities);
         if (request.method === "GET" && url.pathname === "/v2/run-state")
             return json(await storage.transaction((tx) => getAgentSessionRunState(tx, url.searchParams.get("runId") ?? "")));
-        if (request.method === "POST" && url.pathname === "/v2/turns/start") {
-            const input = agentSessionExecutionStartSchema.parse(await readBody(request, storage));
-            return json(await storage.transaction(async (tx) => {
+        if (request.method === "POST" &&
+            (url.pathname === "/v2/turns/start" ||
+                url.pathname === "/v2/turns/prepare")) {
+            const preparing = url.pathname === "/v2/turns/prepare";
+            const body = await readBody(request, storage);
+            const preparation = preparing
+                ? agentSessionExecutionPrepareSchema.parse(body)
+                : undefined;
+            const { resumeWaiting = false, ...input } = preparation ?? {
+                ...agentSessionExecutionStartSchema.parse(body),
+                resumeWaiting: false,
+            };
+            const result = await storage.transaction(async (tx) => {
                 if (await activeComputerOperation(tx))
                     throw new SessionProtocolError("computer_busy");
                 if (await tx.get(key(`outcome:${encodeURIComponent(input.runId)}`)))
@@ -353,6 +364,10 @@ export async function handleAgentSessionExecutionRequest(request, storage) {
                 const previousRevision = await readRevision(tx);
                 const active = await tx.get(key("active-turn"));
                 const resuming = active?.status === "waiting";
+                // Resolve only a waiting continuation under the admission lock. Expired
+                // running attempts still require an explicit expectedAttemptId.
+                const expectedAttemptId = input.expectedAttemptId ??
+                    (resumeWaiting && resuming ? active.attemptId : undefined);
                 if (active) {
                     const sameSelection = active.runId === input.runId &&
                         active.packageVersionId === input.packageVersionId &&
@@ -367,7 +382,7 @@ export async function handleAgentSessionExecutionRequest(request, storage) {
                         throw new SessionProtocolError("turn_already_active");
                     }
                     if (!sameSelection ||
-                        input.expectedAttemptId !== active.attemptId ||
+                        expectedAttemptId !== active.attemptId ||
                         input.requestId === active.requestId)
                         throw new SessionProtocolError("takeover_mismatch");
                     if (!resuming &&
@@ -406,7 +421,8 @@ export async function handleAgentSessionExecutionRequest(request, storage) {
                     });
                 const revision = await advanceRevision(tx);
                 return { duplicate: false, revision, turn };
-            }));
+            });
+            return json(preparing ? { ...result, capabilities } : result);
         }
         if (request.method === "POST" && url.pathname === "/v2/turns/renew") {
             const input = agentSessionExecutionRenewSchema.parse(await readBody(request, storage));
@@ -512,6 +528,10 @@ export async function handleAgentSessionExecutionRequest(request, storage) {
                 const state = await getAgentSessionRunState(tx, input.runId);
                 if (state?.status === "draining")
                     throw new SessionProtocolError("run_cleanup_pending");
+                // An invocation lease may still be live after a child wake claims its
+                // continuation. Only the inbox can seal that logical Run's outcome.
+                if (state?.status === "running" || state?.status === "waiting")
+                    throw new SessionProtocolError("run_not_terminal");
                 const terminalStatus = state?.status === "terminal"
                     ? {
                         completed: "finished",

package/dist/agent-worker.js

@@ -2724,6 +2724,7 @@ export async function runAgentWorkerRequest(input) {
             if (finished)
                 return finished;
             finished = (async () => {
+                const cleanupStartedAt = performance.now();
                 // AI SDK's onEnd may still write context artifacts after a client
                 // disconnects. Settle the model/tools before releasing their filesystem.
                 const results = [];
@@ -2733,10 +2734,13 @@ export async function runAgentWorkerRequest(input) {
                     drained = background.length;
                     results.push(...(await Promise.allSettled(batch)));
                 }
+                const drainedAt = performance.now();
                 if (stopComputer)
                     results.push(...(await Promise.allSettled([stopComputer()])));
+                const computerStoppedAt = performance.now();
                 if (input.onSessionFinished)
                     results.push(...(await Promise.allSettled([input.onSessionFinished()])));
+                const sessionFinishedAt = performance.now();
                 // Keep the Worker lease through draining and Computer cleanup. The
                 // host still owns terminal accounting; hand it a fresh lease only once
                 // this invocation has stopped writing, then stop all Worker heartbeats.
@@ -2756,6 +2760,12 @@ export async function runAgentWorkerRequest(input) {
                         }),
                     ])));
                 }
+                invocationSpan?.setAttributes({
+                    "gea.session.finish.background_ms": Math.round(drainedAt - cleanupStartedAt),
+                    "gea.session.finish.computer_stop_ms": Math.round(computerStoppedAt - drainedAt),
+                    "gea.session.finish.session_release_ms": Math.round(sessionFinishedAt - computerStoppedAt),
+                    "gea.session.finish.lease_handoff_ms": Math.round(performance.now() - sessionFinishedAt),
+                });
                 if (input.inboxExecution) {
                     const status = cancellation.signal.aborted
                         ? "cancelled"
@@ -2848,7 +2858,6 @@ export async function runAgentWorkerRequest(input) {
         let response;
         try {
             if (input.sessionExecution) {
-                await maintainManagedLease();
                 cancellation.signal.throwIfAborted();
                 watchManagedLease();
             }
@@ -2971,6 +2980,10 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
             path: "/v2/turns/renew",
             schema: agentSessionExecutionStartResultSchema,
         });
+        // Renewal atomically validates the attempt and returns cancellation intent;
+        // a separate metadata read before it only adds a network round trip.
+        if (admission.turn.cancelRequestedAt)
+            throw new Error("Agent Session cancelled.");
         mutationSequence = admission.turn.nextMutationSequence;
     }
     const mutateSession = (values) => {
@@ -3508,25 +3521,28 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
             ? (await import("./agent-core.js")).agentCore()
             : input.engine;
     const coreModule = engine?.module();
-    let coreConnection = engine
-        ? await createGeaCoreConnection(host, configuredModel, requestOptions, input.request.signal, input.snapshot.maxOutputTokens)
-        : undefined;
-    const model = engine
-        ? undefined
-        : await createGeaAI(host, configuredModel, requestOptions.extraBody, {
-            kind: "agent-dev",
-            packageName: input.snapshot.name,
-        }, requestOptions.requestBody);
+    const [initialCoreConnection, model, storedContext] = await Promise.all([
+        engine
+            ? createGeaCoreConnection(host, configuredModel, requestOptions, input.request.signal, input.snapshot.maxOutputTokens)
+            : undefined,
+        engine
+            ? undefined
+            : createGeaAI(host, configuredModel, requestOptions.extraBody, {
+                kind: "agent-dev",
+                packageName: input.snapshot.name,
+            }, requestOptions.requestBody),
+        requestSessionJson({
+            chatId: payload.chatId,
+            method: "GET",
+            namespace: input.sessionNamespace,
+            path: sessionExecution ? "/v2/model-context" : "/v1/model-context",
+            schema: agentSessionGetModelContextResultSchema,
+        }),
+    ]);
+    let coreConnection = initialCoreConnection;
     // Model file parts can contain URL and byte values that ordinary JSON cannot preserve.
     const runtimeContext = decodeAgentModelProxyValue(payload.runtimeContext);
     let currentTurn = decodeAgentModelProxyValue(payload.currentTurn);
-    const storedContext = await requestSessionJson({
-        chatId: payload.chatId,
-        method: "GET",
-        namespace: input.sessionNamespace,
-        path: sessionExecution ? "/v2/model-context" : "/v1/model-context",
-        schema: agentSessionGetModelContextResultSchema,
-    });
     if (submittedUser &&
         storedContext.context &&
         storedContext.context.runId !== payload.runId) {
@@ -3814,25 +3830,35 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
         }));
     }
     let mergedResponseMessageCount = 0;
-    const persistModelContext = async () => {
+    let contextCheckpoint;
+    let reusedContextCheckpoints = 0;
+    const persistModelContext = () => {
         const modelContext = {
             contextSize: context.contextSize,
             messages: context.messages.map(encodeAgentModelProxyValue),
-            metadata: context.metadata,
+            metadata: structuredClone(context.metadata),
             runId: payload.runId,
         };
-        if (sessionExecution) {
-            await mutateSession({ context: modelContext });
-            return;
+        const serialized = JSON.stringify(modelContext);
+        // Lifecycle callbacks can checkpoint the same projection several times.
+        // Reuse the pending write too: an unchanged snapshot must still await its
+        // durable commit, including propagating a failed commit to every waiter.
+        if (contextCheckpoint?.serialized === serialized) {
+            persistenceSpan?.setAttribute("gea.session.persistence.context_reused", ++reusedContextCheckpoints);
+            return contextCheckpoint.pending;
         }
-        await requestSessionJson({
-            body: { context: modelContext },
-            chatId: payload.chatId,
-            method: "POST",
-            namespace: input.sessionNamespace,
-            path: "/v1/model-context",
-            schema: agentSessionPutModelContextResultSchema,
-        });
+        const pending = sessionExecution
+            ? mutateSession({ context: modelContext })
+            : requestSessionJson({
+                body: { context: modelContext },
+                chatId: payload.chatId,
+                method: "POST",
+                namespace: input.sessionNamespace,
+                path: "/v1/model-context",
+                schema: agentSessionPutModelContextResultSchema,
+            }).then(() => undefined);
+        contextCheckpoint = { serialized, pending };
+        return pending;
     };
     if (startsNewTurn && findLatestUserMessage(currentTurn)) {
         await persistModelContext();
@@ -4095,9 +4121,16 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
             cursor = page.next_cursor ?? undefined;
         } while (cursor);
     };
+    let finalArtifacts;
+    const refreshFinalArtifacts = () => 
+    // Context callbacks can publish outputs. Both terminal consumers use the
+    // same snapshot, taken only after those callbacks and their work settle.
+    (finalArtifacts ??= contextFinished.then(() => refreshArtifacts()));
     const saveResponseMessage = async ({ responseMessage, }) => {
         const saveStartedAt = performance.now();
-        await refreshArtifacts();
+        const metadata = await contextFinished;
+        const contextFinishedAt = performance.now();
+        await refreshFinalArtifacts();
         const artifactsFinishedAt = performance.now();
         responseMessage.parts = [
             ...responseMessage.parts.filter((part) => part.type !== "data-artifact" ||
@@ -4123,9 +4156,9 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
             context.metadata.suspendedMessageId = responseMessage.id;
         responseMessage.metadata = {
             ...responseMessage.metadata,
-            ...(await contextFinished),
+            ...metadata,
         };
-        const contextFinishedAt = performance.now();
+        const transcriptStartedAt = performance.now();
         // Retain the generated transcript even if context settlement failed.
         // A failed settlement must never produce a successful Run completion.
         if (sessionExecution) {
@@ -4145,9 +4178,9 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
         await persistModelContext();
         const contextPersistedAt = performance.now();
         persistenceSpan?.setAttributes({
-            "gea.session.finish.artifacts_ms": Math.round(artifactsFinishedAt - saveStartedAt),
-            "gea.session.finish.context_wait_ms": Math.round(contextFinishedAt - artifactsFinishedAt),
-            "gea.session.finish.transcript_ms": Math.round(transcriptFinishedAt - contextFinishedAt),
+            "gea.session.finish.artifacts_ms": Math.round(artifactsFinishedAt - contextFinishedAt),
+            "gea.session.finish.context_wait_ms": Math.round(contextFinishedAt - saveStartedAt),
+            "gea.session.finish.transcript_ms": Math.round(transcriptFinishedAt - transcriptStartedAt),
             "gea.session.finish.context_write_ms": Math.round(contextPersistedAt - transcriptFinishedAt),
         });
         const needsInput = !inboxRunner?.suspended &&
@@ -4735,7 +4768,10 @@ async function executeAgentWorkerRequest(input, waitUntil, ownComputer, ownSessi
                 chunk.toolCallId === inboxRunner?.suspendedToolCallId)
                 return;
             if (chunk.type === "finish-step" || chunk.type === "finish") {
-                await refreshArtifacts();
+                if (chunk.type === "finish")
+                    await refreshFinalArtifacts();
+                else
+                    await refreshArtifacts();
                 for (const data of publishedArtifacts.values()) {
                     if (emittedArtifacts.has(data.id))
                         continue;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@gea-ai/agent-sdk",
-  "version": "0.1.260922-alpha.0",
+  "version": "0.1.260922-alpha.2",
   "private": false,
   "homepage": "https://musegea.com/developers",
   "license": "Apache-2.0",
@@ -108,7 +108,7 @@
     "@ai-sdk/otel": "1.0.9",
     "@ai-sdk/provider": "4.0.1",
     "@ai-sdk/provider-utils": "5.0.2",
-    "@gea-ai/contract": "0.1.260922-alpha.0",
+    "@gea-ai/contract": "0.1.260922-alpha.2",
     "@opentelemetry/api": "1.9.1",
     "@opentelemetry/context-async-hooks": "2.10.0",
     "@opentelemetry/core": "2.10.0",