From d0697283d8bd0d6b44dc5074c366ad0a4151f383 Mon Sep 17 00:00:00 2001
From: pc-style <pcstyle@duck.com>
Date: Sun, 17 May 2026 13:44:43 +0200
Subject: [PATCH] Fix Droid review feedback

---
 .../src/provider/Layers/DroidAdapter.test.ts  | 91 +++++++++++++++++++
 .../src/provider/Layers/DroidAdapter.ts       |  2 +-
 .../src/provider/droid/DroidRuntimeEvents.ts  |  5 +-
 3 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/apps/server/src/provider/Layers/DroidAdapter.test.ts b/apps/server/src/provider/Layers/DroidAdapter.test.ts
index cb8c877a71..f03d066b0c 100644
--- a/apps/server/src/provider/Layers/DroidAdapter.test.ts
+++ b/apps/server/src/provider/Layers/DroidAdapter.test.ts
@@ -316,6 +316,97 @@ it.effect("keeps Droid token usage cumulative across turns", () =>
   ).pipe(Effect.provide(testLayer)),
 );
 
+it.effect("accumulates multiple Droid token usage updates within a turn", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const adapter = yield* makeDroidAdapter(settings, {
+        sdk: {
+          createSession: async () =>
+            fakeSession({
+              messages: [
+                {
+                  type: DroidMessageType.TokenUsageUpdate,
+                  inputTokens: 10,
+                  outputTokens: 4,
+                  cacheCreationTokens: 2,
+                  cacheReadTokens: 3,
+                  thinkingTokens: 1,
+                },
+                {
+                  type: DroidMessageType.TokenUsageUpdate,
+                  inputTokens: 5,
+                  outputTokens: 7,
+                  cacheCreationTokens: 0,
+                  cacheReadTokens: 1,
+                  thinkingTokens: 2,
+                },
+                { type: DroidMessageType.TurnComplete, tokenUsage: null },
+              ],
+            }),
+          resumeSession: async () => fakeSession({}),
+        },
+      });
+      const eventsFiber = yield* adapter.streamEvents.pipe(
+        Stream.filter((event) => event.threadId === threadId),
+        Stream.take(6),
+        Stream.runCollect,
+        Effect.forkChild,
+      );
+
+      yield* adapter.startSession({
+        threadId,
+        provider: ProviderDriverKind.make("droid"),
+        runtimeMode: "full-access",
+      });
+      yield* adapter.sendTurn({ threadId, input: "hello" });
+
+      const events = Array.from(yield* Fiber.join(eventsFiber).pipe(Effect.timeout("2 seconds")));
+      const usageEvents = events.filter((event) => event.type === "thread.token-usage.updated");
+
+      assert.deepEqual(
+        usageEvents.map((event) =>
+          event.type === "thread.token-usage.updated" ? event.payload.usage : undefined,
+        ),
+        [
+          {
+            usedTokens: 20,
+            inputTokens: 15,
+            cachedInputTokens: 3,
+            outputTokens: 5,
+            reasoningOutputTokens: 1,
+            lastUsedTokens: 20,
+            lastInputTokens: 15,
+            lastCachedInputTokens: 3,
+            lastOutputTokens: 5,
+            lastReasoningOutputTokens: 1,
+          },
+          {
+            usedTokens: 35,
+            inputTokens: 21,
+            cachedInputTokens: 4,
+            outputTokens: 14,
+            reasoningOutputTokens: 3,
+            lastUsedTokens: 15,
+            lastInputTokens: 6,
+            lastCachedInputTokens: 1,
+            lastOutputTokens: 9,
+            lastReasoningOutputTokens: 2,
+          },
+        ],
+      );
+      const completed = events.find((event) => event.type === "turn.completed");
+      assert.equal(completed?.type, "turn.completed");
+      if (completed?.type === "turn.completed") {
+        assert.equal(completed.payload.state, "completed");
+        assert.equal(
+          (completed.payload as { usage?: { usedTokens?: number } }).usage?.usedTokens,
+          35,
+        );
+      }
+    }),
+  ).pipe(Effect.provide(testLayer)),
+);
+
 it.effect("maps Droid medium access to medium autonomy", () =>
   Effect.scoped(
     Effect.gen(function* () {
diff --git a/apps/server/src/provider/Layers/DroidAdapter.ts b/apps/server/src/provider/Layers/DroidAdapter.ts
index 1a779f03c3..faba84f883 100644
--- a/apps/server/src/provider/Layers/DroidAdapter.ts
+++ b/apps/server/src/provider/Layers/DroidAdapter.ts
@@ -425,7 +425,7 @@ export function makeDroidAdapter(settings: DroidSettings, options?: DroidAdapter
           ...eventBase(context),
           type: "session.exited",
           payload: { reason: "Session stopped", recoverable: false, exitKind: "graceful" },
-        });
+        }).pipe(Effect.ignore);
       });
 
     return {
diff --git a/apps/server/src/provider/droid/DroidRuntimeEvents.ts b/apps/server/src/provider/droid/DroidRuntimeEvents.ts
index fb9f32358f..ba7a494ced 100644
--- a/apps/server/src/provider/droid/DroidRuntimeEvents.ts
+++ b/apps/server/src/provider/droid/DroidRuntimeEvents.ts
@@ -186,7 +186,10 @@ export async function handleDroidMessage(input: {
         },
       });
     case DroidMessageType.TokenUsageUpdate:
-      context.activeTokenUsage = toTokenUsageSnapshot(message, context.activeTokenUsageBaseline);
+      context.activeTokenUsage = toTokenUsageSnapshot(
+        message,
+        context.activeTokenUsage ?? context.activeTokenUsageBaseline,
+      );
       context.cumulativeTokenUsage = context.activeTokenUsage;
       return emitNow({
         ...base(),