@@ -3,6 +3,7 @@ package agent
33import (
44 "context"
55 _ "embed"
6+ "encoding/json"
67 "errors"
78 "fmt"
89 "log/slog"
@@ -177,7 +178,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
177178 // Before each step create the new assistant message
178179 PrepareStep: func(callContext context.Context, options ai.PrepareStepFunctionOptions) (_ context.Context, prepared ai.PrepareStepResult, err error) {
179180 var assistantMsg message.Message
180− assistantMsg, err = a.messages.Create(genCtx, call.SessionID, message.CreateMessageParams{
181+ assistantMsg, err = a.messages.Create(callContext, call.SessionID, message.CreateMessageParams{
181182 Role: message.Assistant,
182183 Parts: []message.ContentPart{},
183184 Model: a.largeModel.ModelCfg.Model,
@@ -187,7 +188,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
187188 return callContext, prepared, err
188189 }
189190
190− callContext = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
191+ callContext = context.WithValue(callContext, tools.MessageIDContextKey, assistantMsg.ID)
191192
192193 currentAssistant = &assistantMsg
193194
@@ -200,7 +201,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
200201 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
201202 a.messageQueue.Del(call.SessionID)
202203 for _, queued := range queuedCalls {
203− userMessage, createErr := a.createUserMessage(genCtx, queued)
204+ userMessage, createErr := a.createUserMessage(callContext, queued)
204205 if createErr != nil {
205206 return callContext, prepared, createErr
206207 }
@@ -291,12 +292,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
291292 IsError: isError,
292293 Metadata: result.ClientMetadata,
293294 }
294− a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
295+ _, err := a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
295296 Role: message.Tool,
296297 Parts: []message.ContentPart{
297298 toolResult,
298299 },
299300 })
301+ if err != nil {
302+ return err
303+ }
300304 return a.messages.Update(genCtx, *currentAssistant)
301305 },
302306 OnStepFinish: func(stepResult ai.StepResult) error {
@@ -328,16 +332,29 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
328332 }
329333 toolCalls := currentAssistant.ToolCalls()
330334 toolResults := currentAssistant.ToolResults()
335+ // INFO: we use the parent context here because the genCtx has been cancelled
336+ msgs, createErr := a.messages.List(ctx, currentSession.ID)
337+ if createErr != nil {
338+ return nil, createErr
339+ }
331340 for _, tc := range toolCalls {
332341 if !tc.Finished {
333342 tc.Finished = true
334343 tc.Input = "{}"
335344 }
336345 currentAssistant.AddToolCall(tc)
346+
337347 found := false
338− for _, tr := range toolResults {
339− if tr.ToolCallID == tc.ID {
340− found = true
348+ for _, msg := range msgs {
349+ if msg.Role == message.Tool {
350+ for _, tr := range toolResults {
351+ if tr.ToolCallID == tc.ID {
352+ found = true
353+ break
354+ }
355+ }
356+ }
357+ if found {
341358 break
342359 }
343360 }
@@ -348,12 +365,21 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
348365 } else if isPermissionErr {
349366 content = "Permission denied"
350367 }
351− currentAssistant.AddToolResult(message.ToolResult{
368+ toolResult := message.ToolResult{
352369 ToolCallID: tc.ID,
353370 Name: tc.Name,
354371 Content: content,
355372 IsError: true,
373+ }
374+ _, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
375+ Role: message.Tool,
376+ Parts: []message.ContentPart{
377+ toolResult,
378+ },
356379 })
380+ if createErr != nil {
381+ return nil, createErr
382+ }
357383 }
358384 }
359385 if isCancelErr {
@@ -363,13 +389,11 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
363389 } else {
364390 currentAssistant.AddFinish(message.FinishReasonError, "API Error", err.Error())
365391 }
366− // INFO: we use the parent context here because the genCtx might have been cancelled
392+ // INFO: we use the parent context here because the genCtx has been cancelled
367393 updateErr := a.messages.Update(ctx, *currentAssistant)
368394 if updateErr != nil {
369395 return nil, updateErr
370396 }
371− }
372− if err != nil {
373397 return nil, err
374398 }
375399 wg.Wait()
@@ -541,22 +565,32 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
541565 return
542566 }
543567
568+ var maxOutput int64 = 40
569+ if a.smallModel.CatwalkCfg.CanReason {
570+ maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
571+ }
572+
544573 agent := ai.NewAgent(a.smallModel.Model,
545− ai.WithSystemPrompt(string(titlePrompt)),
546− ai.WithMaxOutputTokens(40),
574+ ai.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
575+ ai.WithMaxOutputTokens(maxOutput),
547576 )
548577
549578 resp, err := agent.Stream(ctx, ai.AgentStreamCall{
550− Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s", prompt),
579+ Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", prompt),
551580 })
552581 if err != nil {
553582 slog.Error("error generating title", "err", err)
554583 return
555584 }
556585
586+ data, _ := json.Marshal(resp)
587+
588+ slog.Info("Title Response")
589+ slog.Info(string(data))
557590 title := resp.Response.Content.Text()
558591
559592 title = strings.ReplaceAll(title, "\n", " ")
593+ slog.Info(title)
560594
561595 // remove thinking tags if present
562596 if idx := strings.Index(title, "</think>"); idx > 0 {