Skip to content

Commit b5006b8

Browse files
brandonkachencodebuff-team
brandonkachen
and
codebuff-team
committed
Improve type safety for ContentBlock state management
Refactor message state updates to use explicit ContentBlock types: - Replace inline object literals with typed ContentBlock variables - Add const assertions for status fields ('running', 'complete') - Ensure blocks arrays are properly typed as ContentBlock[] - Make agent block initialization more explicit with empty arrays - Improve readability and maintainability of complex state updates 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent a2c8894 commit b5006b8

File tree

1 file changed

+181
-127
lines changed

1 file changed

+181
-127
lines changed

‎cli/src/hooks/use-send-message.ts‎

Lines changed: 181 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -162,11 +162,13 @@ export const useSendMessage = ({
162162
if (msg.id === aiMessageId && msg.blocks) {
163163
const newBlocks = msg.blocks.map((block) => {
164164
if (block.type === 'agent' && block.agentId === agentId) {
165-
const agentBlocks = block.blocks || []
165+
const agentBlocks: ContentBlock[] = block.blocks
166+
? [...block.blocks]
167+
: []
166168
if (update.type === 'text' && update.content) {
167169
const lastBlock = agentBlocks[agentBlocks.length - 1]
168170
if (lastBlock && lastBlock.type === 'text') {
169-
const updatedLastBlock = {
171+
const updatedLastBlock: ContentBlock = {
170172
...lastBlock,
171173
content: lastBlock.content + update.content,
172174
}
@@ -234,27 +236,38 @@ export const useSendMessage = ({
234236
})
235237
setMessages((prev) =>
236238
prev.map((msg) => {
237-
if (msg.id === aiMessageId) {
238-
const blocks = msg.blocks || []
239-
const lastBlock = blocks[blocks.length - 1]
239+
if (msg.id !== aiMessageId) {
240+
return msg
241+
}
240242

241-
if (lastBlock && lastBlock.type === 'text') {
242-
const newContent = lastBlock.content + text
243-
return {
244-
...msg,
245-
blocks: [
246-
...blocks.slice(0, -1),
247-
{ type: 'text', content: newContent },
248-
],
249-
}
250-
} else {
251-
return {
252-
...msg,
253-
blocks: [...blocks, { type: 'text', content: text }],
254-
}
243+
const blocks: ContentBlock[] = msg.blocks
244+
? [...msg.blocks]
245+
: []
246+
const lastBlock = blocks[blocks.length - 1]
247+
248+
if (lastBlock && lastBlock.type === 'text') {
249+
const newContent = lastBlock.content + text
250+
const updatedTextBlock: ContentBlock = {
251+
type: 'text',
252+
content: newContent,
253+
}
254+
return {
255+
...msg,
256+
blocks: [
257+
...blocks.slice(0, -1),
258+
updatedTextBlock,
259+
],
255260
}
256261
}
257-
return msg
262+
263+
const newTextBlock: ContentBlock = {
264+
type: 'text',
265+
content: text,
266+
}
267+
return {
268+
...msg,
269+
blocks: [...blocks, newTextBlock],
270+
}
258271
}),
259272
)
260273
},
@@ -366,26 +379,37 @@ export const useSendMessage = ({
366379
})
367380
setMessages((prev) =>
368381
prev.map((msg) => {
369-
if (msg.id === aiMessageId) {
370-
const blocks = msg.blocks || []
371-
const lastBlock = blocks[blocks.length - 1]
382+
if (msg.id !== aiMessageId) {
383+
return msg
384+
}
372385

373-
if (lastBlock && lastBlock.type === 'text') {
374-
return {
375-
...msg,
376-
blocks: [
377-
...blocks.slice(0, -1),
378-
{ type: 'text', content: lastBlock.content + text },
379-
],
380-
}
381-
} else {
382-
return {
383-
...msg,
384-
blocks: [...blocks, { type: 'text', content: text }],
385-
}
386+
const blocks: ContentBlock[] = msg.blocks
387+
? [...msg.blocks]
388+
: []
389+
const lastBlock = blocks[blocks.length - 1]
390+
391+
if (lastBlock && lastBlock.type === 'text') {
392+
const updatedTextBlock: ContentBlock = {
393+
type: 'text',
394+
content: lastBlock.content + text,
395+
}
396+
return {
397+
...msg,
398+
blocks: [
399+
...blocks.slice(0, -1),
400+
updatedTextBlock,
401+
],
386402
}
387403
}
388-
return msg
404+
405+
const newTextBlock: ContentBlock = {
406+
type: 'text',
407+
content: text,
408+
}
409+
return {
410+
...msg,
411+
blocks: [...blocks, newTextBlock],
412+
}
389413
}),
390414
)
391415
return
@@ -474,25 +498,28 @@ export const useSendMessage = ({
474498
)
475499
setMessages((prev) =>
476500
prev.map((msg) => {
477-
if (msg.id === aiMessageId) {
478-
const blocks = msg.blocks || []
479-
return {
480-
...msg,
481-
blocks: [
482-
...blocks,
483-
{
484-
type: 'agent',
485-
agentId: event.agentId,
486-
agentName: event.agentType || 'Agent',
487-
agentType: event.agentType || 'unknown',
488-
content: '',
489-
status: 'running',
490-
initialPrompt: '',
491-
},
492-
],
493-
}
501+
if (msg.id !== aiMessageId) {
502+
return msg
503+
}
504+
505+
const blocks: ContentBlock[] = msg.blocks
506+
? [...msg.blocks]
507+
: []
508+
const newAgentBlock: ContentBlock = {
509+
type: 'agent',
510+
agentId: event.agentId,
511+
agentName: event.agentType || 'Agent',
512+
agentType: event.agentType || 'unknown',
513+
content: '',
514+
status: 'running' as const,
515+
blocks: [] as ContentBlock[],
516+
initialPrompt: '',
517+
}
518+
519+
return {
520+
...msg,
521+
blocks: [...blocks, newAgentBlock],
494522
}
495-
return msg
496523
}),
497524
)
498525

@@ -555,27 +582,31 @@ export const useSendMessage = ({
555582

556583
setMessages((prev) =>
557584
prev.map((msg) => {
558-
if (msg.id === aiMessageId) {
559-
const blocks = msg.blocks || []
560-
const newAgentBlocks = agents.map(
561-
(agent: any, index: number) => ({
562-
type: 'agent' as const,
563-
agentId: `${toolCallId}-${index}`,
564-
agentName: agent.agent_type || 'Agent',
565-
agentType: agent.agent_type || 'unknown',
566-
content: agent.prompt || '',
567-
status: 'running' as const,
568-
blocks: [],
569-
initialPrompt: agent.prompt || '',
570-
}),
571-
)
585+
if (msg.id !== aiMessageId) {
586+
return msg
587+
}
572588

573-
return {
574-
...msg,
575-
blocks: [...blocks, ...newAgentBlocks],
576-
}
589+
const existingBlocks: ContentBlock[] = msg.blocks
590+
? [...msg.blocks]
591+
: []
592+
593+
const newAgentBlocks: ContentBlock[] = agents.map(
594+
(agent: any, index: number) => ({
595+
type: 'agent',
596+
agentId: `${toolCallId}-${index}`,
597+
agentName: agent.agent_type || 'Agent',
598+
agentType: agent.agent_type || 'unknown',
599+
content: agent.prompt || '',
600+
status: 'running' as const,
601+
blocks: [] as ContentBlock[],
602+
initialPrompt: agent.prompt || '',
603+
}),
604+
)
605+
606+
return {
607+
...msg,
608+
blocks: [...existingBlocks, ...newAgentBlocks],
577609
}
578-
return msg
579610
}),
580611
)
581612

@@ -613,40 +644,58 @@ export const useSendMessage = ({
613644

614645
setMessages((prev) =>
615646
prev.map((msg) => {
616-
if (msg.id === aiMessageId && msg.blocks) {
617-
const blocks = msg.blocks.map((block) => {
618-
if (block.type === 'agent' && block.agentId === agentId) {
619-
const agentBlocks = block.blocks || []
620-
return {
621-
...block,
622-
blocks: [
623-
...agentBlocks,
624-
{ type: 'tool', toolCallId, toolName, input },
625-
],
626-
}
627-
}
628-
return block
629-
})
630-
return { ...msg, blocks }
647+
if (msg.id !== aiMessageId || !msg.blocks) {
648+
return msg
631649
}
632-
return msg
650+
651+
const updatedBlocks: ContentBlock[] = msg.blocks.map(
652+
(block) => {
653+
if (block.type !== 'agent' || block.agentId !== agentId) {
654+
return block
655+
}
656+
657+
const agentBlocks: ContentBlock[] = block.blocks
658+
? [...block.blocks]
659+
: []
660+
const newToolBlock: ContentBlock = {
661+
type: 'tool',
662+
toolCallId,
663+
toolName,
664+
input,
665+
}
666+
667+
return {
668+
...block,
669+
blocks: [...agentBlocks, newToolBlock],
670+
}
671+
},
672+
)
673+
674+
return { ...msg, blocks: updatedBlocks }
633675
}),
634676
)
635677
} else {
636678
// Top-level tool call
637679
setMessages((prev) =>
638680
prev.map((msg) => {
639-
if (msg.id === aiMessageId) {
640-
const blocks = msg.blocks || []
641-
return {
642-
...msg,
643-
blocks: [
644-
...blocks,
645-
{ type: 'tool', toolCallId, toolName, input },
646-
],
647-
}
681+
if (msg.id !== aiMessageId) {
682+
return msg
683+
}
684+
685+
const existingBlocks: ContentBlock[] = msg.blocks
686+
? [...msg.blocks]
687+
: []
688+
const newToolBlock: ContentBlock = {
689+
type: 'tool',
690+
toolCallId,
691+
toolName,
692+
input,
693+
}
694+
695+
return {
696+
...msg,
697+
blocks: [...existingBlocks, newToolBlock],
648698
}
649-
return msg
650699
}),
651700
)
652701
}
@@ -702,9 +751,13 @@ export const useSendMessage = ({
702751
contentPreview: content.substring(0, 100),
703752
})
704753

754+
const resultTextBlock: ContentBlock = {
755+
type: 'text',
756+
content,
757+
}
705758
return {
706759
...block,
707-
blocks: [{ type: 'text', content }],
760+
blocks: [resultTextBlock],
708761
status: 'complete' as const,
709762
}
710763
}
@@ -817,35 +870,36 @@ export const useSendMessage = ({
817870
if (isAborted) {
818871
setMessages((prev) =>
819872
prev.map((msg) => {
820-
if (msg.id === aiMessageId) {
821-
const blocks = msg.blocks || []
822-
const lastBlock = blocks[blocks.length - 1]
873+
if (msg.id !== aiMessageId) {
874+
return msg
875+
}
823876

824-
if (lastBlock && lastBlock.type === 'text') {
825-
return {
826-
...msg,
827-
blocks: [
828-
...blocks.slice(0, -1),
829-
{
830-
type: 'text',
831-
content:
832-
lastBlock.content + '\n\n[response interrupted]',
833-
},
834-
],
835-
isComplete: true,
836-
}
837-
} else {
838-
return {
839-
...msg,
840-
blocks: [
841-
...blocks,
842-
{ type: 'text', content: '[response interrupted]' },
843-
],
844-
isComplete: true,
845-
}
877+
const blocks: ContentBlock[] = msg.blocks
878+
? [...msg.blocks]
879+
: []
880+
const lastBlock = blocks[blocks.length - 1]
881+
882+
if (lastBlock && lastBlock.type === 'text') {
883+
const interruptedBlock: ContentBlock = {
884+
type: 'text',
885+
content: `${lastBlock.content}\n\n[response interrupted]`,
886+
}
887+
return {
888+
...msg,
889+
blocks: [...blocks.slice(0, -1), interruptedBlock],
890+
isComplete: true,
846891
}
847892
}
848-
return msg
893+
894+
const interruptionNotice: ContentBlock = {
895+
type: 'text',
896+
content: '[response interrupted]',
897+
}
898+
return {
899+
...msg,
900+
blocks: [...blocks, interruptionNotice],
901+
isComplete: true,
902+
}
849903
}),
850904
)
851905
} else {

0 commit comments

Comments
 (0)