refactor: address code review — extract stream helpers, improve type narrowing

Agent-Logs-Url: https://github.com/objectstack-ai/spec/sessions/7b5017d9-7b75-41e6-a0ad-6eb250b8d770

Co-authored-by: hotlong <50353452+hotlong@users.noreply.github.com>

Author: Copilot

packages/services/service-ai/src/__tests__/chatbot-features.test.ts

@@ -259,7 +259,13 @@ describe('AIService.chatWithTools', () => {
     // The error message should be in the tool result
     const secondCallMessages = (adapter.chat as any).mock.calls[1][0] as ModelMessage[];
     const toolMsg = secondCallMessages.find(m => m.role === 'tool');
-    const toolContent = Array.isArray(toolMsg?.content) ? (toolMsg.content[0] as any)?.output?.value : toolMsg?.content;
+    let toolContent: string | undefined;
+    if (toolMsg?.role === 'tool' && Array.isArray(toolMsg.content)) {
+      const firstResult = toolMsg.content[0];
+      if ('output' in firstResult && firstResult.output && typeof firstResult.output === 'object' && 'value' in firstResult.output) {
+        toolContent = String(firstResult.output.value);
+      }
+    }
     expect(toolContent).toContain('Tool crashed');
   });
 

packages/services/service-ai/src/__tests__/objectql-conversation-service.test.ts

@@ -265,8 +265,12 @@ describe('ObjectQLConversationService', () => {
 
     const updated = await service.addMessage(conv.id, msg);
     expect(updated.messages).toHaveLength(1);
-    const toolContent = updated.messages[0].content as any[];
-    expect(toolContent[0].toolCallId).toBe('call_abc');
+    const firstMsg = updated.messages[0];
+    if (firstMsg.role === 'tool' && Array.isArray(firstMsg.content)) {
+      expect(firstMsg.content[0].toolCallId).toBe('call_abc');
+    } else {
+      throw new Error('Expected tool message with array content');
+    }
   });
 
   it('should add an assistant message with toolCalls', async () => {
@@ -280,10 +284,14 @@ describe('ObjectQLConversationService', () => {
 
     const updated = await service.addMessage(conv.id, msg);
     expect(updated.messages).toHaveLength(1);
-    const assistantContent = updated.messages[0].content as any[];
-    const toolCallParts = assistantContent.filter((p: any) => p.type === 'tool-call');
-    expect(toolCallParts).toHaveLength(1);
-    expect(toolCallParts[0].toolName).toBe('get_weather');
+    const firstMsg = updated.messages[0];
+    if (firstMsg.role === 'assistant' && Array.isArray(firstMsg.content)) {
+      const toolCallParts = firstMsg.content.filter((p) => p.type === 'tool-call');
+      expect(toolCallParts).toHaveLength(1);
+      expect(toolCallParts[0].toolName).toBe('get_weather');
+    } else {
+      throw new Error('Expected assistant message with array content');
+    }
   });
 
   it('should throw when adding message to non-existent conversation', async () => {

packages/services/service-ai/src/ai-service.ts

@@ -19,6 +19,25 @@ import { ToolRegistry } from './tools/tool-registry.js';
 import type { ToolExecutionResult } from './tools/tool-registry.js';
 import { InMemoryConversationService } from './conversation/in-memory-conversation-service.js';
 
+// ── Stream event helpers ──────────────────────────────────────────
+// These helpers construct properly-typed Vercel AI SDK stream parts
+// to avoid repeated `as unknown as TextStreamPart<ToolSet>` casts.
+
+/** Create a text-delta stream part. */
+function textDeltaPart(id: string, text: string): TextStreamPart<ToolSet> {
+  return { type: 'text-delta', id, text } as TextStreamPart<ToolSet>;
+}
+
+/** Create a finish stream part from an AIResult. */
+function finishPart(result?: AIResult): TextStreamPart<ToolSet> {
+  return {
+    type: 'finish',
+    finishReason: 'stop',
+    totalUsage: result?.usage ?? { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
+    rawFinishReason: 'stop',
+  } as unknown as TextStreamPart<ToolSet>;
+}
+
 /**
  * Configuration for AIService.
  */
@@ -92,13 +111,8 @@ export class AIService implements IAIService {
     if (!this.adapter.streamChat) {
       // Fallback: emit the entire response as a single text-delta + finish
       const result = await this.adapter.chat(messages, options);
-      yield { type: 'text-delta', id: 'fallback', text: result.content } as TextStreamPart<ToolSet>;
-      yield {
-        type: 'finish',
-        finishReason: 'stop' as const,
-        totalUsage: result.usage ?? { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
-        rawFinishReason: 'stop',
-      } as unknown as TextStreamPart<ToolSet>;
+      yield textDeltaPart('fallback', result.content);
+      yield finishPart(result);
       return;
     }
 
@@ -287,13 +301,8 @@ export class AIService implements IAIService {
 
       if (!result.toolCalls || result.toolCalls.length === 0) {
         // Final round — return the probed result without an extra model call
-        yield { type: 'text-delta', id: 'stream', text: result.content } as TextStreamPart<ToolSet>;
-        yield {
-          type: 'finish',
-          finishReason: 'stop' as const,
-          totalUsage: result.usage ?? { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
-          rawFinishReason: 'stop',
-        } as unknown as TextStreamPart<ToolSet>;
+        yield textDeltaPart('stream', result.content);
+        yield finishPart(result);
         return;
       }
 
@@ -342,12 +351,7 @@ export class AIService implements IAIService {
     }
     const finalOptions = { ...chatOptions, tools: undefined, toolChoice: undefined };
     const result = await this.adapter.chat(conversation, finalOptions);
-    yield { type: 'text-delta', id: 'stream', text: result.content } as TextStreamPart<ToolSet>;
-    yield {
-      type: 'finish',
-      finishReason: 'stop' as const,
-      totalUsage: result.usage ?? { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
-      rawFinishReason: 'stop',
-    } as unknown as TextStreamPart<ToolSet>;
+    yield textDeltaPart('stream', result.content);
+    yield finishPart(result);
   }
 }

packages/services/service-ai/src/conversation/objectql-conversation-service.ts

@@ -183,7 +183,7 @@ export class ObjectQLConversationService implements IAIConversationService {
         contentStr = message.content;
       } else {
         const parts = message.content;
-        const textParts = parts.filter(p => p.type === 'text').map(p => 'text' in p ? p.text : '');
+        const textParts = parts.filter((p): p is { type: 'text'; text: string } => p.type === 'text').map(p => p.text);
         const toolCalls = parts.filter(p => p.type === 'tool-call');
         contentStr = textParts.join('');
         if (toolCalls.length > 0) toolCallsJson = JSON.stringify(toolCalls);
@@ -280,7 +280,7 @@ export class ObjectQLConversationService implements IAIConversationService {
       }
       case 'tool': {
         const toolResults = this.safeParse<ToolResultPart[]>(row.content);
-        if (toolResults && Array.isArray(toolResults) && toolResults.length > 0 && toolResults[0]?.type === 'tool-result') {
+        if (toolResults && toolResults.length > 0 && toolResults[0]?.type === 'tool-result') {
           return { role: 'tool', content: toolResults };
         }
         // Backward compat: old format was a plain string