feat: streamed usage on finish chunk and pageUrl skill arming in sessions

docs/LOCAL-ENGINE-API.md

@@ -91,7 +91,8 @@ keywords:
   so `홈택스` matches `홈택스에서`.
 - **urls** match `redrob_context.page_url`. The scheme is ignored, `*` stays inside one
   path segment, `**` crosses segments, and a trailing `/*` in `urls:` covers the whole
-  site including its root.
+  site including its root. In agent sessions the same matcher reads the optional
+  `pageUrl` field of the session prompt body.
 - Armed bodies go in as one extra `system` message after the caller's own system
   messages, labelled with what armed each one. A skill with neither list never arms
   and behaves as before; in agent sessions an armable skill is left out of the standing

packages/redrob/src/session/prompt.ts

@@ -698,6 +698,7 @@ const layer = Layer.effect(
         },
         system: input.system,
         format: input.format,
+        ...(input.pageUrl?.trim() ? { pageUrl: input.pageUrl.trim() } : {}),
       }
 
       const current = yield* sessions.get(input.sessionID).pipe(Effect.orDie)
@@ -1300,18 +1301,18 @@ const layer = Layer.effect(
              * otherwise arm on half of what the user said.
              *
              * The URL is deliberately NOT fetched. Arming by URL glob needs the active
-             * tab, and the only way to get it is a browser request -- a round trip on
-             * every single message, which with no extension attached spends the request
-             * deadline before the model sees anything. A twenty second pause before each
-             * reply is a worse product than URL arming is a better one. Keywords arm
-             * today; the URL half arms once a session carries its tab URL as state the
-             * engine already holds, rather than as a question it has to ask.
+             * tab, and asking the browser for it is a round trip on every single message,
+             * which with no extension attached spends the request deadline before the
+             * model sees anything. Instead the caller states it: `pageUrl` on the prompt
+             * is stored on the user message, and arms through the same matcher the
+             * chat-completions route runs on `redrob_context.page_url`.
              */
             const armedSkills = yield* sys.armedSkills(agent, {
               prompt: (lastUserMsg?.parts ?? [])
                 .filter((part) => part.type === "text")
                 .map((part) => (part.type === "text" ? part.text : ""))
                 .join("\n"),
+              url: lastUser.pageUrl,
             })
             const system = [
               ...env,
@@ -1563,6 +1564,10 @@ export const PromptInput = Schema.Struct({
   format: Schema.optional(SessionV1.Format),
   system: Schema.optional(Schema.String),
   variant: Schema.optional(Schema.String),
+  pageUrl: Schema.optional(Schema.String).annotate({
+    description:
+      "URL of the caller's active browser tab. Skills whose `urls:` globs match it arm for this turn, as on /v1/chat/completions.",
+  }),
   parts: Schema.Array(
     Schema.Union([
       SessionV1.TextPartInput,

packages/schema/src/v1/session.ts

@@ -351,6 +351,12 @@ export const User = Schema.Struct({
   }),
   system: Schema.optional(Schema.String),
   tools: Schema.optional(Schema.Record(Schema.String, Schema.Boolean)),
+  /**
+   * The URL of the browser tab the caller was on when it sent this message. Skills whose
+   * `urls:` globs match it arm for the turn, the same matcher `/v1/chat/completions` runs
+   * on `redrob_context.page_url`. Absent for callers with no tab.
+   */
+  pageUrl: Schema.optional(Schema.String),
 }).annotate({ identifier: "UserMessage" })
 export type User = Types.DeepMutable<Schema.Schema.Type<typeof User>>
 

packages/sdk/js/src/v2/gen/sdk.gen.ts

@@ -3871,6 +3871,7 @@ export class Session2 extends HeyApiClient {
       format?: OutputFormat
       system?: string
       variant?: string
+      pageUrl?: string
       parts?: Array<TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput>
     },
     options?: Options<never, ThrowOnError>,
@@ -3891,6 +3892,7 @@ export class Session2 extends HeyApiClient {
             { in: "body", key: "format" },
             { in: "body", key: "system" },
             { in: "body", key: "variant" },
+            { in: "body", key: "pageUrl" },
             { in: "body", key: "parts" },
           ],
         },
@@ -4224,6 +4226,7 @@ export class Session2 extends HeyApiClient {
       format?: OutputFormat
       system?: string
       variant?: string
+      pageUrl?: string
       parts?: Array<TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput>
     },
     options?: Options<never, ThrowOnError>,
@@ -4244,6 +4247,7 @@ export class Session2 extends HeyApiClient {
             { in: "body", key: "format" },
             { in: "body", key: "system" },
             { in: "body", key: "variant" },
+            { in: "body", key: "pageUrl" },
             { in: "body", key: "parts" },
           ],
         },

packages/sdk/js/src/v2/gen/types.gen.ts

@@ -310,6 +310,7 @@ export type UserMessage = {
   tools?: {
     [key: string]: boolean
   }
+  pageUrl?: string
 }
 
 export type ProviderAuthError = {
@@ -10384,6 +10385,7 @@ export type SessionPromptData = {
     format?: OutputFormat
     system?: string
     variant?: string
+    pageUrl?: string
     parts: Array<TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput>
   }
   path: {
@@ -10731,6 +10733,7 @@ export type SessionPromptAsyncData = {
     format?: OutputFormat
     system?: string
     variant?: string
+    pageUrl?: string
     parts: Array<TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput>
   }
   path: {

packages/server/src/handlers/chat-completion.ts

@@ -342,20 +342,27 @@ function toJsonBody(id: string, model: string, response: LLMResponse): unknown {
         finish_reason: toFinishReason(response.finishReason),
       },
     ],
-    ...(usage
-      ? {
-          usage: {
-            prompt_tokens: usage.inputTokens ?? 0,
-            completion_tokens: usage.outputTokens ?? 0,
-            total_tokens: usage.totalTokens ?? 0,
-          },
-        }
-      : {}),
+    ...(usage ? { usage: toWireUsage(usage) } : {}),
+  }
+}
+
+/** The engine's usage, in OpenAI's `usage` shape. One mapping for the JSON body and the stream. */
+function toWireUsage(usage: NonNullable<LLMResponse["usage"]>) {
+  return {
+    prompt_tokens: usage.inputTokens ?? 0,
+    completion_tokens: usage.outputTokens ?? 0,
+    total_tokens: usage.totalTokens ?? 0,
   }
 }
 
 /** One SSE frame carrying an OpenAI chunk. */
-function chunkFrame(id: string, model: string, delta: unknown, finishReason: string | null): Sse.Event {
+function chunkFrame(
+  id: string,
+  model: string,
+  delta: unknown,
+  finishReason: string | null,
+  usage?: ReturnType<typeof toWireUsage>,
+): Sse.Event {
   return {
     _tag: "Event",
     event: "message",
@@ -366,6 +373,7 @@ function chunkFrame(id: string, model: string, delta: unknown, finishReason: str
       created: seconds(),
       model,
       choices: [{ index: 0, delta, finish_reason: finishReason }],
+      ...(usage ? { usage } : {}),
     }),
   }
 }
@@ -386,13 +394,25 @@ function doneFrame(): Sse.Event {
  * caller would read a truncated answer as a complete one.
  */
 function toChunks(id: string, model: string, events: Stream.Stream<LLMEvent, LLMError>) {
+  // Usage rides on the finish frame, as OpenAI's own final chunk carries it. A provider that
+  // reports it only per step still gets counted: the last step's usage stands in. Without this the
+  // streaming path dropped usage entirely, and the browser's memory runs recorded 0 tokens.
+  let stepUsage: LLMResponse["usage"]
   return events.pipe(
-    Stream.flatMap((event: LLMEvent) => Stream.fromIterable(eventFrames(id, model, event))),
+    Stream.flatMap((event: LLMEvent) => {
+      if (event.type === "step-finish" && event.usage !== undefined) stepUsage = event.usage
+      return Stream.fromIterable(eventFrames(id, model, event, stepUsage))
+    }),
     Stream.concat(Stream.make(doneFrame())),
   )
 }
 
-function eventFrames(id: string, model: string, event: LLMEvent): ReadonlyArray<Sse.Event> {
+function eventFrames(
+  id: string,
+  model: string,
+  event: LLMEvent,
+  stepUsage?: LLMResponse["usage"],
+): ReadonlyArray<Sse.Event> {
   switch (event.type) {
     case "text-delta":
       return [chunkFrame(id, model, { content: event.text }, null)]
@@ -402,8 +422,10 @@ function eventFrames(id: string, model: string, event: LLMEvent): ReadonlyArray<
       const [call] = toWireToolCalls([{ id: event.id, name: event.name, input: event.input }])
       return [chunkFrame(id, model, { tool_calls: [{ index: 0, ...call }] }, null)]
     }
-    case "finish":
-      return [chunkFrame(id, model, {}, toFinishReason(event.reason))]
+    case "finish": {
+      const usage = event.usage ?? stepUsage
+      return [chunkFrame(id, model, {}, toFinishReason(event.reason), usage ? toWireUsage(usage) : undefined)]
+    }
     case "provider-error":
       // Surfaced as data, not as a stream failure: headers are already sent, so
       // this is the only way the caller learns the turn broke.