@@ -3,6 +3,7 @@ package agent
@@ -177,7 +178,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
177 // Before each step create the new assistant message
178 // Before each step create the new assistant message
178 PrepareStep: func(callContext context.Context, options ai.PrepareStepFunctionOptions) (_ context.Context, prepared ai.PrepareStepResult, err error) {
179 PrepareStep: func(callContext context.Context, options ai.PrepareStepFunctionOptions) (_ context.Context, prepared ai.PrepareStepResult, err error) {
179 var assistantMsg message.Message
180 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{
181 Role: message.Assistant,
182 Role: message.Assistant,
182 Parts: []message.ContentPart{},
183 Parts: []message.ContentPart{},
183 Model: a.largeModel.ModelCfg.Model,
184 Model: a.largeModel.ModelCfg.Model,
@@ -187,7 +188,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
187 return callContext, prepared, err
188 return callContext, prepared, err
190− callContext = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
191+ callContext = context.WithValue(callContext, tools.MessageIDContextKey, assistantMsg.ID)
192 currentAssistant = &assistantMsg
193 currentAssistant = &assistantMsg
@@ -200,7 +201,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
200 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
201 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
201 a.messageQueue.Del(call.SessionID)
202 a.messageQueue.Del(call.SessionID)
202 for _, queued := range queuedCalls {
203 for _, queued := range queuedCalls {
203− userMessage, createErr := a.createUserMessage(genCtx, queued)
204+ userMessage, createErr := a.createUserMessage(callContext, queued)
204 if createErr != nil {
205 if createErr != nil {
205 return callContext, prepared, createErr
206 return callContext, prepared, createErr
@@ -291,12 +292,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
291 IsError: isError,
292 IsError: isError,
292 Metadata: result.ClientMetadata,
293 Metadata: result.ClientMetadata,
294− a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
295+ _, err := a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
295 Role: message.Tool,
296 Role: message.Tool,
296 Parts: []message.ContentPart{
297 Parts: []message.ContentPart{
297 toolResult,
298 toolResult,
300 return a.messages.Update(genCtx, *currentAssistant)
304 return a.messages.Update(genCtx, *currentAssistant)
302 OnStepFinish: func(stepResult ai.StepResult) error {
306 OnStepFinish: func(stepResult ai.StepResult) error {
@@ -328,16 +332,29 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
329 toolCalls := currentAssistant.ToolCalls()
333 toolCalls := currentAssistant.ToolCalls()
330 toolResults := currentAssistant.ToolResults()
334 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
331 for _, tc := range toolCalls {
340 for _, tc := range toolCalls {
332 if !tc.Finished {
341 if !tc.Finished {
333 tc.Finished = true
342 tc.Finished = true
334 tc.Input = "{}"
343 tc.Input = "{}"
336 currentAssistant.AddToolCall(tc)
345 currentAssistant.AddToolCall(tc)
337 found := false
347 found := false
348+ for _, msg := range msgs {
349+ if msg.Role == message.Tool {
338− for _, tr := range toolResults {
350+ for _, tr := range toolResults {
339− if tr.ToolCallID == tc.ID {
351+ if tr.ToolCallID == tc.ID {
340− found = true
352+ found = true
@@ -348,12 +365,21 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
348 } else if isPermissionErr {
365 } else if isPermissionErr {
349 content = "Permission denied"
366 content = "Permission denied"
351− currentAssistant.AddToolResult(message.ToolResult{
368+ toolResult := message.ToolResult{
352 ToolCallID: tc.ID,
369 ToolCallID: tc.ID,
353 Name: tc.Name,
370 Name: tc.Name,
354 Content: content,
371 Content: content,
355 IsError: true,
372 IsError: true,
374+ _, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
376+ Parts: []message.ContentPart{
380+ if createErr != nil {
381+ return nil, createErr
359 if isCancelErr {
385 if isCancelErr {
@@ -363,13 +389,11 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
364 currentAssistant.AddFinish(message.FinishReasonError, "API Error", err.Error())
390 currentAssistant.AddFinish(message.FinishReasonError, "API Error", err.Error())
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
367 updateErr := a.messages.Update(ctx, *currentAssistant)
393 updateErr := a.messages.Update(ctx, *currentAssistant)
368 if updateErr != nil {
394 if updateErr != nil {
369 return nil, updateErr
395 return nil, updateErr
373 return nil, err
397 return nil, err
375 wg.Wait()
399 wg.Wait()
@@ -541,22 +565,32 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
568+ var maxOutput int64 = 40
569+ if a.smallModel.CatwalkCfg.CanReason {
570+ maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
544 agent := ai.NewAgent(a.smallModel.Model,
573 agent := ai.NewAgent(a.smallModel.Model,
545− ai.WithSystemPrompt(string(titlePrompt)),
574+ ai.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
546− ai.WithMaxOutputTokens(40),
575+ ai.WithMaxOutputTokens(maxOutput),
549 resp, err := agent.Stream(ctx, ai.AgentStreamCall{
578 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),
552 if err != nil {
581 if err != nil {
553 slog.Error("error generating title", "err", err)
582 slog.Error("error generating title", "err", err)
586+ data, _ := json.Marshal(resp)
588+ slog.Info("Title Response")
589+ slog.Info(string(data))
557 title := resp.Response.Content.Text()
590 title := resp.Response.Content.Text()
559 title = strings.ReplaceAll(title, "\n", " ")
592 title = strings.ReplaceAll(title, "\n", " ")
561 // remove thinking tags if present
595 // remove thinking tags if present
562 if idx := strings.Index(title, "</think>"); idx > 0 {
596 if idx := strings.Index(title, "</think>"); idx > 0 {