Server tool results missing from tool-call output field in UIMessage

[flybayer]

TanStack AI version

0.2.0

Framework/Library version

react 18.2.0

Describe the bug and the steps to reproduce it

Summary: Server-side tool results don't update the tool-call part's output field, causing inconsistent UI state between server and client tool execution.

Expected behavior: Both server-executed tools and client-executed tools should have their results visible on the tool-call part's output field, and the state should transition to 'complete'.

Actual behavior:

• Client tools: tool-call part shows state: 'complete' with output populated
• Server tools: tool-call part shows state: 'input-complete' with no output (result only exists in a separate tool-result part)

Root cause: In StreamProcessor, handleToolResultChunk() only creates a tool-result part via updateToolResultPart(), but doesn't update the corresponding tool-call part's output field or state. In contrast, addToolResult() (used for client tools) does both:

// addToolResult() does this (lines 283-300):
let updatedMessages = updateToolCallWithOutput(  // Updates tool-call output
  this.messages,
  toolCallId,
  output,
  error ? 'input-complete' : undefined,
  error,
)
updatedMessages = updateToolResultPart(...)  // Also creates tool-result part

// handleToolResultChunk() only does this (lines 691-699):
this.messages = updateToolResultPart(...)  // Only creates tool-result part
// Missing: updateToolCallWithOutput() call

Steps to reproduce:

1. Define a tool with toolDefinition() that has both inputSchema and outputSchema
2. Create a server-side implementation using .server() with an execute function
3. Create a client-side implementation using .client() with an execute function
4. Call both tools in a chat session
5. Observe the UIMessage.parts - client tool-call parts have output and state: 'complete', server tool-call parts only have state: 'input-complete' with no output

Example to reproduce

import { toolDefinition, chat } from '@tanstack/ai'
import { z } from 'zod'

// Shared tool definition
const myToolDef = toolDefinition({
  name: 'myTool',
  description: 'A test tool',
  inputSchema: z.object({ input: z.string() }),
  outputSchema: z.object({ result: z.string() }),
})

// Server tool - has execute, runs on server
const serverTool = myToolDef.server(async (args) => {
  return { result: `Server processed: ${args.input}` }
})

// Client tool - has execute, runs on client
const clientTool = myToolDef.client(async (args) => {
  return { result: `Client processed: ${args.input}` }
})

// When server tool executes:
// - StreamProcessor receives tool_result chunk
// - handleToolResultChunk() creates tool-result part
// - tool-call part remains at state: 'input-complete', output: undefined ❌

// When client tool executes:
// - StreamProcessor receives tool-input-available chunk
// - Client executes tool and calls addToolResult()
// - addToolResult() updates tool-call part with output AND creates tool-result part
// - tool-call part shows state: 'complete', output: {...} ✅

Screenshots or Videos (Optional)

In the screenshot:

• getCurrentPipelineConfig (client tool): Shows "complete" status with RESULT section visible
• getPipelineSchema (server tool): Shows "input-complete" status with no RESULT section
• validatePipelineConfig (server tool): Shows "input-complete" status with no RESULT section
• applyPipelineConfig (client tool): Shows "complete" status

Suggested Fix

In processor.ts, handleToolResultChunk() should also update the tool-call part similar to addToolResult():

private handleToolResultChunk(
  chunk: Extract<StreamChunk, { type: 'tool_result' }>,
): void {
  const state: ToolResultState = 'complete'

  // Emit legacy handler
  this.handlers.onToolResultStateChange?.(
    chunk.toolCallId,
    chunk.content,
    state,
  )

  // Update UIMessage if we have a current assistant message
  if (this.currentAssistantMessageId) {
    // NEW: Also update the tool-call part's output field
    const output = this.parseToolResultContent(chunk.content)
    this.messages = updateToolCallWithOutput(
      this.messages,
      chunk.toolCallId,
      output,
    )
    
    // Existing: Create tool-result part
    this.messages = updateToolResultPart(
      this.messages,
      this.currentAssistantMessageId,
      chunk.toolCallId,
      chunk.content,
      state,
    )
    this.emitMessagesChange()
  }
}

Do you intend to try to help solve this bug with your own PR?

None

Terms & Code of Conduct

• I agree to follow this project's Code of Conduct
• I understand that if my bug cannot be reliable reproduced in a debuggable environment, it will probably not be fixed and this issue may even be closed.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🔗 Related PRs

#141 - Add TanStack Start server code example for chat API [merged]

👤 Suggested Assignees

• harlley

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!