@anvia/studio 1.0.1 → 1.0.2

package/README.md

@@ -35,7 +35,7 @@ const agent = new Agent({
   instructions: "Answer support questions clearly.",
 });
 
-new Studio([agent]).start({
+await new Studio([agent]).serve({
   port: 4021,
 });
 ```
@@ -46,6 +46,25 @@ Then open:
 http://localhost:4021/ui/playground
 ```
 
+## Graceful shutdown
+
+`serve()` handles both `SIGINT` and `SIGTERM`. Studio stops accepting work, aborts active Agent and
+Pipeline runs, waits for their cancellation observers, and then runs `onShutdown`. Use that callback
+to close observability clients or other caller-owned resources:
+
+```ts
+await new Studio([agent]).serve({
+  port: 4021,
+  shutdownTimeoutMs: 30_000,
+  onShutdown: async () => {
+    await Promise.all([lens.close(), langfuse.close(), otelSdk.shutdown()]);
+  },
+});
+```
+
+`shutdown()` provides the same draining behavior for application-managed lifecycles. `close()`
+remains synchronous for compatibility: it aborts active work but does not wait for cleanup.
+
 ## Multi-Provider Models
 
 Studio can expose a shared model catalog and let each agent choose from registered providers:

package/dist/index.d.ts

@@ -581,7 +581,7 @@ type StudioSessionStore = StudioMemoryStore & {
     listSessionLogs?(options: StudioSessionLogListOptions): StudioSessionLogEntry[] | Promise<StudioSessionLogEntry[]>;
     deleteSession?(id: string): boolean | Promise<boolean>;
 };
-type StudioTraceStatus = "running" | "success" | "suspended" | "error";
+type StudioTraceStatus = "running" | "success" | "suspended" | "error" | "cancelled";
 type StudioTraceObservationKind = "agent" | "generation" | "tool";
 type StudioTraceObservation = {
     id: string;
@@ -848,6 +848,7 @@ type StudioServeOptions = {
     hostname?: string;
     log?: boolean;
     handleSignals?: boolean;
+    shutdownTimeoutMs?: number;
 };
 type StudioServeLifecycleOptions = Omit<StudioServeOptions, "handleSignals"> & {
     signal?: AbortSignal;
@@ -1011,7 +1012,7 @@ type AgentRunRequest = (AgentRunRequestBase & {
 });
 type AgentRunResponse = Exclude<AgentOutcome, AgentInteractionOutcome> | Omit<AgentInteractionOutcome, "continuation" | "messages">;
 type AgentRunStreamEvent = AgentStreamEvent | StudioSessionLogEvent | StudioPipelineLogEvent | StudioPipelineFinalEvent;
-type StudioErrorCode = "bad_request" | "conflict" | "forbidden" | "not_found" | "payload_too_large" | "unsupported_capability" | "internal_error";
+type StudioErrorCode = "bad_request" | "conflict" | "forbidden" | "not_found" | "payload_too_large" | "service_unavailable" | "unsupported_capability" | "internal_error";
 type StudioErrorResponse = {
     error: {
         code: StudioErrorCode;
@@ -1025,6 +1026,9 @@ type AnviaStudio = {
     fetch(request: Request): Response | Promise<Response>;
     config(): StudioConfig;
     close(): void;
+    shutdown(options?: {
+        timeoutMs?: number;
+    }): Promise<void>;
 };
 
 type StudioTraceObserverOptions = {
@@ -1042,7 +1046,8 @@ declare class Studio implements AnviaStudio {
     private studio;
     private server;
     private websocketServer;
-    private sigintHandler;
+    private signalHandlers;
+    private shutdownPromise;
     constructor(targets?: StudioTarget[], options?: StudioOptions);
     get app(): Hono;
     fetch(request: Request): Response | Promise<Response>;
@@ -1051,6 +1056,12 @@ declare class Studio implements AnviaStudio {
     start(serveOptions?: StudioServeOptions): this;
     serve(serveOptions?: StudioServeLifecycleOptions): Promise<void>;
     close(): void;
+    shutdown(options?: {
+        timeoutMs?: number;
+    }): Promise<void>;
+    private shutdownResources;
+    private removeSignalHandlers;
+    private closeNetwork;
     private logAddress;
 }
 

package/dist/index.js

@@ -5,6 +5,9 @@ import { serve } from "@hono/node-server";
 import { Hono as HonoApp } from "hono";
 import { WebSocketServer } from "ws";
 
+// src/traces/trace-observer.ts
+import { AgentRunCancelledError } from "@anvia/core/agent";
+
 // src/runtime/json.ts
 function toJsonValue(value) {
   return toJsonValueInternal(value, /* @__PURE__ */ new WeakSet());
@@ -213,7 +216,7 @@ var StudioRunTraceObserver = class {
     });
   }
   async error(args) {
-    await this.save("error", {
+    await this.save(args.status === "cancelled" ? "cancelled" : "error", {
       endedAt: /* @__PURE__ */ new Date(),
       error: serializeUnknown(args.error),
       usage: args.usage,
@@ -355,11 +358,12 @@ var ChildAgentToolTraceAccumulator = class {
     if (child.type === "error") {
       const metadata = this.childMetadata(agentId, agentName, childTurn);
       if (isRecord(child.usage)) metadata.usage = toJsonValue(child.usage);
+      const status = child.error instanceof AgentRunCancelledError ? "cancelled" : "error";
       this.completedObservations.push(
         traceObservation({
           kind: "tool",
           name: `${agentLabel(agentId, agentName)}.error`,
-          status: "error",
+          status,
           turn: this.parent.turn,
           startedAt: /* @__PURE__ */ new Date(),
           error: serializeUnknown(child.error),
@@ -385,7 +389,7 @@ var ChildAgentToolTraceAccumulator = class {
         parentObservationId,
         kind: "agent",
         name: `${agentLabel(agentStart.agentId, agentStart.agentName)}.run`,
-        status: agentChildren.some((observation) => observation.status === "error") ? "error" : "success",
+        status: agentChildren.some((observation) => observation.status === "error") ? "error" : agentChildren.some((observation) => observation.status === "cancelled") ? "cancelled" : "success",
         turn: this.parent.turn,
         startedAt,
         endedAt,
@@ -3089,16 +3093,30 @@ function registerAgentRunRoute(app, props) {
   app.post("/agents/:agentId/runs", (c) => handleAgentRun(c, props));
 }
 async function handleAgentRun(c, props) {
-  const prepared = await prepareAgentRun(c, props);
-  if (prepared instanceof Response) {
-    return prepared;
+  const lease = props.runLifecycle.start(c.req.raw.signal);
+  if (lease === void 0) {
+    return errorResponse(c, 503, "service_unavailable", "Anvia Studio is shutting down");
   }
-  if (prepared.body.stream === true) {
-    return handleStreamingAgentRun(c, prepared, props);
+  try {
+    const prepared = await prepareAgentRun(c, props, lease.abortSignal);
+    if (prepared instanceof Response) {
+      lease.finish();
+      return prepared;
+    }
+    if (prepared.body.stream === true) {
+      return handleStreamingAgentRun(c, prepared, props, lease);
+    }
+    try {
+      return await handleBufferedAgentRun(c, prepared, props);
+    } finally {
+      lease.finish();
+    }
+  } catch (error) {
+    lease.finish();
+    throw error;
   }
-  return handleBufferedAgentRun(c, prepared, props);
 }
-async function prepareAgentRun(c, props) {
+async function prepareAgentRun(c, props, abortSignal) {
   const agentId = c.req.param("agentId");
   const agent = props.agentMap.get(agentId);
   if (agent === void 0) {
@@ -3132,7 +3150,7 @@ async function prepareAgentRun(c, props) {
     if (body.stream !== void 0) resumedBody = { ...resumedBody, stream: body.stream };
     if (body.metadata !== void 0) resumedBody = { ...resumedBody, metadata: body.metadata };
     if (body.trace !== void 0) resumedBody = { ...resumedBody, trace: body.trace };
-    const runOptions = createRunOptions(resumedBody, agentId, source.session, c.req.raw.signal);
+    const runOptions = createRunOptions(resumedBody, agentId, source.session, abortSignal);
     const execution2 = {
       generate: (options) => source.runAgent.generate({
         continuation: claimed.continuation,
@@ -3211,7 +3229,7 @@ async function prepareAgentRun(c, props) {
     execution,
     failureMessages: void 0,
     memoryCompactionLogged: false,
-    options: createRunOptions(body, agentId, session, c.req.raw.signal),
+    options: createRunOptions(body, agentId, session, abortSignal),
     runAgent,
     runId,
     runStartedAt,
@@ -3339,7 +3357,7 @@ function createRunOptions(body, agentId, session, abortSignal) {
   }
   return options;
 }
-function handleStreamingAgentRun(c, run, props) {
+function handleStreamingAgentRun(c, run, props, lease) {
   const streamOptions = withInternalAgentRunOptions({ ...run.options }, { runId: run.runId });
   const runStream = registerStreamingContinuation(
     run.execution.stream(streamOptions),
@@ -3367,7 +3385,7 @@ function handleStreamingAgentRun(c, run, props) {
     });
     stream = persistedRun.events;
   }
-  return streamAgentRunEvents(c, stream, {
+  return streamAgentRunEvents(c, finishRunLease(stream, lease), {
     runId: run.runId,
     onCancel: async () => {
       const persistence = [];
@@ -3386,6 +3404,13 @@ function handleStreamingAgentRun(c, run, props) {
     }
   });
 }
+async function* finishRunLease(events, lease) {
+  try {
+    yield* events;
+  } finally {
+    lease.finish();
+  }
+}
 async function handleBufferedAgentRun(c, run, props) {
   try {
     const runtimeOptions = await startBufferedSessionRun(run);
@@ -3680,7 +3705,7 @@ function parseTraceStatus(value) {
   if (status === void 0) {
     return void 0;
   }
-  return status === "running" || status === "success" || status === "suspended" || status === "error" ? status : false;
+  return status === "running" || status === "success" || status === "suspended" || status === "error" || status === "cancelled" ? status : false;
 }
 function parseAfter(value) {
   if (value === void 0 || value.trim().length === 0) {
@@ -4997,7 +5022,7 @@ function registerPipelineRoutes(app, props) {
     if ("error" in body) {
       return body.error;
     }
-    return executePipelineRun(c, props, pipeline, body);
+    return executeTrackedPipelineRun(c, props, pipeline, body);
   });
   app.post("/pipelines/:pipelineId/runs/:runId/replay", async (c) => {
     const pipeline = props.pipelineMap.get(c.req.param("pipelineId"));
@@ -5033,10 +5058,22 @@ function registerPipelineRoutes(app, props) {
       metadata: replayMetadata(sourceRun.metadata, body.metadata, sourceRun.runId)
     };
     if (body.stream !== void 0) replayRequest.stream = body.stream;
-    return executePipelineRun(c, props, pipeline, replayRequest);
+    return executeTrackedPipelineRun(c, props, pipeline, replayRequest);
   });
 }
-async function executePipelineRun(c, props, pipeline, body) {
+async function executeTrackedPipelineRun(c, props, pipeline, body) {
+  const lease = props.runLifecycle.start(c.req.raw.signal);
+  if (lease === void 0) {
+    return errorResponse(c, 503, "service_unavailable", "Anvia Studio is shutting down");
+  }
+  try {
+    return await executePipelineRun(c, props, pipeline, body, lease);
+  } catch (error) {
+    lease.finish();
+    throw error;
+  }
+}
+async function executePipelineRun(c, props, pipeline, body, lease) {
   const runId = globalThis.crypto.randomUUID();
   const startedAt = Date.now();
   const startedAtIso = new Date(startedAt).toISOString();
@@ -5063,12 +5100,13 @@ async function executePipelineRun(c, props, pipeline, body) {
       runId,
       input: body.input,
       startedAt,
-      startedAtIso
+      startedAtIso,
+      onFinish: lease.finish
     };
     if (body.metadata !== void 0) streamOptions.metadata = body.metadata;
     if (props.logStore !== void 0) streamOptions.logStore = props.logStore;
     if (props.runStore !== void 0) streamOptions.runStore = props.runStore;
-    streamOptions.abortSignal = c.req.raw.signal;
+    streamOptions.abortSignal = lease.abortSignal;
     return streamPipelineRun(c, streamOptions);
   }
   try {
@@ -5077,7 +5115,7 @@ async function executePipelineRun(c, props, pipeline, body) {
       input: body.input,
       runId,
       metadata: body.metadata,
-      abortSignal: c.req.raw.signal,
+      abortSignal: lease.abortSignal,
       observer: {
         async onEvent(event) {
           await appendPipelineLog(props.logStore, pipelineStageLog(event, pipeline.id));
@@ -5134,6 +5172,8 @@ async function executePipelineRun(c, props, pipeline, body) {
       pipelineRunFailedLog(pipeline.id, runId, error, startedAt)
     );
     return errorResponse(c, 500, "internal_error", "Pipeline run failed", serializeError(error));
+  } finally {
+    lease.finish();
   }
 }
 function pipelineDetail(pipeline) {
@@ -5232,7 +5272,11 @@ async function* pipelineRunEvents(props) {
       yield next.value;
     }
   } finally {
-    await run;
+    try {
+      await run;
+    } finally {
+      props.onFinish();
+    }
   }
 }
 async function savePipelineRun(store, input) {
@@ -5320,6 +5364,61 @@ function parsePipelineLogAfter(value) {
   return after;
 }
 
+// src/runtime/run-lifecycle.ts
+var StudioRunLifecycle = class {
+  shutdownController = new AbortController();
+  active = /* @__PURE__ */ new Set();
+  drainPromise;
+  resolveDrain;
+  closed = false;
+  start(requestSignal) {
+    if (this.closed) return void 0;
+    const token = /* @__PURE__ */ Symbol("studio-run");
+    this.active.add(token);
+    let finished = false;
+    return {
+      abortSignal: AbortSignal.any([requestSignal, this.shutdownController.signal]),
+      finish: () => {
+        if (finished) return;
+        finished = true;
+        this.active.delete(token);
+        if (this.active.size === 0) {
+          this.resolveDrain?.();
+          this.resolveDrain = void 0;
+          this.drainPromise = void 0;
+        }
+      }
+    };
+  }
+  close() {
+    if (this.closed) return;
+    this.closed = true;
+    this.shutdownController.abort(new Error("Anvia Studio is shutting down."));
+  }
+  async drain(timeoutMs) {
+    this.close();
+    if (this.active.size === 0) return;
+    this.drainPromise ??= new Promise((resolve2) => {
+      this.resolveDrain = resolve2;
+    });
+    await withTimeout(this.drainPromise, timeoutMs);
+  }
+};
+async function withTimeout(promise, timeoutMs) {
+  let timer;
+  const timeout = new Promise((_resolve, reject) => {
+    timer = setTimeout(
+      () => reject(new Error(`Anvia Studio shutdown timed out after ${timeoutMs}ms.`)),
+      timeoutMs
+    );
+  });
+  try {
+    await Promise.race([promise, timeout]);
+  } finally {
+    if (timer !== void 0) clearTimeout(timer);
+  }
+}
+
 // src/runtime/sandbox-views.ts
 import { upgradeWebSocket } from "@hono/node-server";
 import { getConnInfo } from "@hono/node-server/conninfo";
@@ -6810,7 +6909,8 @@ var Studio = class {
   studio;
   server;
   websocketServer;
-  sigintHandler;
+  signalHandlers;
+  shutdownPromise;
   constructor(targets = [], options = {}) {
     this.options = studioOptionsFromTargets(targets, options);
     this.studio = createStudioApp(this.options);
@@ -6832,6 +6932,7 @@ var Studio = class {
   start(serveOptions = {}) {
     this.close();
     this.studio = createStudioApp(this.options);
+    this.shutdownPromise = void 0;
     const port = serveOptions.port ?? Number(process.env.RUNNER_PORT ?? 4021);
     const serverOptions = {
       fetch: (request, env) => this.studio.app.fetch(request, env),
@@ -6851,16 +6952,27 @@ var Studio = class {
       this.logAddress(host, port);
     }
     if (serveOptions.handleSignals ?? true) {
-      this.sigintHandler = () => {
-        this.close();
-        process.exit(0);
+      const shutdownFor = (signal) => {
+        process.exitCode = signal === "SIGINT" ? 130 : 143;
+        void this.shutdown(shutdownOptions(serveOptions.shutdownTimeoutMs)).catch(
+          (error) => {
+            console.error(error);
+            process.exitCode = 1;
+          }
+        );
+      };
+      this.signalHandlers = {
+        sigint: () => shutdownFor("SIGINT"),
+        sigterm: () => shutdownFor("SIGTERM")
       };
-      process.once("SIGINT", this.sigintHandler);
+      process.once("SIGINT", this.signalHandlers.sigint);
+      process.once("SIGTERM", this.signalHandlers.sigterm);
     }
     return this;
   }
   async serve(serveOptions = {}) {
-    const { onShutdown, signal, ...startOptions } = serveOptions;
+    const { onShutdown, shutdownTimeoutMs, signal, ...startOptions } = serveOptions;
+    const failures = [];
     try {
       this.start({ ...startOptions, log: false, handleSignals: false });
       const server = this.server;
@@ -6875,21 +6987,55 @@ var Studio = class {
         this.logAddress(startOptions.hostname ?? "localhost", port);
       }
       await waitForShutdown(signal);
-    } finally {
-      this.close();
+    } catch (error) {
+      failures.push(error);
+    }
+    try {
+      await this.shutdown(shutdownOptions(shutdownTimeoutMs));
+    } catch (error) {
+      failures.push(error);
+    }
+    try {
       await onShutdown?.();
+    } catch (error) {
+      failures.push(error);
+    }
+    if (failures.length === 1) throw failures[0];
+    if (failures.length > 1) {
+      throw new AggregateError(failures, "Failed to run and shut down Anvia Studio.");
     }
   }
   close() {
-    if (this.sigintHandler !== void 0) {
-      process.off("SIGINT", this.sigintHandler);
-      this.sigintHandler = void 0;
+    this.removeSignalHandlers();
+    this.closeNetwork();
+    this.studio.close();
+  }
+  shutdown(options = {}) {
+    const timeoutMs = options.timeoutMs ?? 3e4;
+    if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
+      return Promise.reject(
+        new TypeError("Anvia Studio shutdown timeoutMs must be a positive number.")
+      );
     }
+    this.shutdownPromise ??= this.shutdownResources(timeoutMs);
+    return this.shutdownPromise;
+  }
+  async shutdownResources(timeoutMs) {
+    this.removeSignalHandlers();
+    this.closeNetwork();
+    await this.studio.shutdown({ timeoutMs });
+  }
+  removeSignalHandlers() {
+    if (this.signalHandlers === void 0) return;
+    process.off("SIGINT", this.signalHandlers.sigint);
+    process.off("SIGTERM", this.signalHandlers.sigterm);
+    this.signalHandlers = void 0;
+  }
+  closeNetwork() {
     this.server?.close();
     this.server = void 0;
     this.websocketServer?.close();
     this.websocketServer = void 0;
-    this.studio.close();
   }
   logAddress(host, port) {
     if (isStudioUiEnabled(this.options.ui)) {
@@ -6900,6 +7046,9 @@ var Studio = class {
     }
   }
 };
+function shutdownOptions(timeoutMs) {
+  return timeoutMs === void 0 ? {} : { timeoutMs };
+}
 function waitForServerListening(server) {
   if (server.listening) return Promise.resolve();
   return new Promise((resolve2, reject) => {
@@ -7045,6 +7194,7 @@ function createStudioApp(options) {
   const pipelineMap = new Map(pipelines.map((pipeline) => [pipeline.id, pipeline]));
   const evalMap = new Map(options.evals.map((suite) => [suite.id ?? suite.name, suite]));
   const continuationRegistry = createStudioContinuationRegistry();
+  const runLifecycle = new StudioRunLifecycle();
   const sandboxRegistry = createStudioSandboxRegistry(agents, options.sandboxes ?? []);
   const memorySources = createStudioMemorySourceRegistry(agents, stores.sessions);
   const app = new HonoApp();
@@ -7107,7 +7257,8 @@ function createStudioApp(options) {
   registerKnowledgeRoutes(app, knowledgeOptions);
   const pipelineOptions = {
     pipelines,
-    pipelineMap
+    pipelineMap,
+    runLifecycle
   };
   if (stores.pipelineLogs !== void 0) pipelineOptions.logStore = stores.pipelineLogs;
   if (stores.pipelineRuns !== void 0) pipelineOptions.runStore = stores.pipelineRuns;
@@ -7116,7 +7267,8 @@ function createStudioApp(options) {
     agentMap,
     stores,
     modelRegistry,
-    continuationRegistry
+    continuationRegistry,
+    runLifecycle
   });
   if (memorySources.size > 0 || stores.sessions !== void 0) {
     registerMemoryRoutes(app, {
@@ -7139,6 +7291,14 @@ function createStudioApp(options) {
     app.all(`/${capability}`, (c) => unsupportedCapability(c, capability));
     app.all(`/${capability}/*`, (c) => unsupportedCapability(c, capability));
   }
+  let closing = false;
+  const beginClose = () => {
+    if (closing) return;
+    closing = true;
+    continuationRegistry.clear();
+    closeSandboxViews();
+    runLifecycle.close();
+  };
   const studio = {
     app,
     fetch(request) {
@@ -7148,8 +7308,11 @@ function createStudioApp(options) {
       return buildConfig(options, agents, pipelines, stores, sandboxRegistry.size);
     },
     close() {
-      continuationRegistry.clear();
-      closeSandboxViews();
+      beginClose();
+    },
+    async shutdown(shutdownOptions2 = {}) {
+      beginClose();
+      await runLifecycle.drain(shutdownOptions2.timeoutMs ?? 3e4);
     }
   };
   if (stores.sessions !== void 0) Object.assign(studio, { sessionStore: stores.sessions });