Ver Fonte

fix(core): settle malformed tool input on failure (#36477)

Kit Langton há 1 mês atrás
pai
commit
8e657c7db5

+ 24 - 14
packages/core/src/session/runner/publish-llm-event.ts

@@ -220,8 +220,31 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
     yield* flushFragments()
   })
 
+  const failTools = Effect.fnUntraced(function* (error: SessionError.Error, mode: "all" | "hosted" | "uncalled") {
+    let failed = false
+    for (const [callID, tool] of tools) {
+      if (
+        tool.settled ||
+        (mode === "hosted" && !tool.providerExecuted) ||
+        (mode === "uncalled" && tool.called)
+      )
+        continue
+      tool.settled = true
+      failed = true
+      yield* events.publish(SessionEvent.Tool.Failed, {
+        sessionID: input.sessionID,
+        assistantMessageID: tool.assistantMessageID,
+        callID,
+        error,
+        executed: tool.providerExecuted,
+      })
+    }
+    return failed
+  })
+
   const failAssistant = Effect.fnUntraced(function* (error: SessionError.Error, replace = false) {
     yield* flush()
+    yield* failTools(error, "uncalled")
     yield* startAssistant()
     if (replace || stepFailure === undefined) stepFailure = error
   })
@@ -245,20 +268,7 @@ export const createLLMEventPublisher = (events: Pick<EventV2.Interface, "publish
     error: SessionError.Error,
     hostedOnly = false,
   ) {
-    let failed = false
-    for (const [callID, tool] of tools) {
-      if (tool.settled || (hostedOnly && !tool.providerExecuted)) continue
-      tool.settled = true
-      failed = true
-      yield* events.publish(SessionEvent.Tool.Failed, {
-        sessionID: input.sessionID,
-        assistantMessageID: tool.assistantMessageID,
-        callID,
-        error,
-        executed: tool.providerExecuted,
-      })
-    }
-    return failed
+    return yield* failTools(error, hostedOnly ? "hosted" : "all")
   })
 
   const assistantMessageIDForTool = (callID: string) => {

+ 52 - 1
packages/core/test/session-runner.test.ts

@@ -6,6 +6,7 @@ import {
   Model,
   ToolFailure,
   TransportReason,
+  InvalidProviderOutputReason,
   InvalidRequestReason,
   RateLimitReason,
   type LLMClientShape,
@@ -716,7 +717,18 @@ const verifyPartialFlushOnFailure = (kind: FragmentKind) =>
         type: "assistant",
         finish: "error",
         error: { type: "provider.transport", message: "Provider unavailable" },
-        content: [fixture.expectedContent],
+        content: [
+          kind === "tool input"
+            ? {
+                type: "tool",
+                id: fragmentID(kind, "partial"),
+                state: {
+                  status: "error",
+                  error: { type: "provider.transport", message: "Provider unavailable" },
+                },
+              }
+            : fixture.expectedContent,
+        ],
       },
     ])
     expect(requests).toHaveLength(1)
@@ -3876,6 +3888,45 @@ describe("SessionRunnerLLM", () => {
     }),
   )
 
+  it.effect("settles malformed streamed tool input before the provider failure", () =>
+    Effect.gen(function* () {
+      const session = yield* setup
+      yield* admit(session, "Call a malformed tool")
+      const failure = new LLMError({
+        module: "test",
+        method: "stream",
+        reason: new InvalidProviderOutputReason({ message: "Invalid JSON input for tool call echo" }),
+      })
+      responseStream = Stream.fromIterable([
+        LLMEvent.stepStart({ index: 0 }),
+        LLMEvent.toolInputStart({ id: "call-malformed", name: "echo" }),
+        LLMEvent.toolInputDelta({ id: "call-malformed", name: "echo", text: '{"text":"partial' }),
+      ]).pipe(Stream.concat(Stream.fail(failure)))
+
+      expect(yield* session.resume(sessionID).pipe(Effect.flip)).toBe(failure)
+      const assistant = requireAssistant(yield* session.context(sessionID))
+
+      response = reply.stop()
+      yield* admit(session, "Continue")
+      yield* session.resume(sessionID)
+
+      expect(yield* recordedStepSettlementEvents(sessionID, assistant.id)).toMatchObject([
+        { type: "session.step.started.1" },
+        {
+          type: "session.tool.failed.1",
+          data: {
+            callID: "call-malformed",
+            error: { type: "provider.invalid-output", message: "Invalid JSON input for tool call echo" },
+          },
+        },
+        {
+          type: "session.step.failed.1",
+          data: { error: { type: "provider.invalid-output", message: "Invalid JSON input for tool call echo" } },
+        },
+      ])
+    }),
+  )
+
   it.effect("does not continue automatically after a provider error follows a local tool call", () =>
     Effect.gen(function* () {
       const session = yield* setup