implement nested tool calls and initial setup for result metadata

3 files changed, 31 insertions, 6 deletions

diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 998dc1551..b01ffec3c 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -305,6 +305,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
 		assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
 			Role:  message.Assistant,
 			Parts: []message.ContentPart{},
+			Model: c.model.ID,
 		})
 		if err != nil {
 			return err
diff --git a/internal/llm/tools/bash.go b/internal/llm/tools/bash.go
index 4e80ae60a..d20afb7f2 100644
--- a/internal/llm/tools/bash.go
+++ b/internal/llm/tools/bash.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/kujtimiihoxha/termai/internal/config"
 	"github.com/kujtimiihoxha/termai/internal/llm/tools/shell"
@@ -21,6 +22,9 @@ type BashPermissionsParams struct {
 	Timeout int    `json:"timeout"`
 }
 
+type BashToolResponseMetadata struct {
+	Took int64 `json:"took"`
+}
 type bashTool struct {
 	permissions permission.Service
 }
@@ -272,11 +276,13 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 			return NewTextErrorResponse("permission denied"), nil
 		}
 	}
+	startTime := time.Now()
 	shell := shell.GetPersistentShell(config.WorkingDirectory())
 	stdout, stderr, exitCode, interrupted, err := shell.Exec(ctx, params.Command, params.Timeout)
 	if err != nil {
 		return NewTextErrorResponse(fmt.Sprintf("error executing command: %s", err)), nil
 	}
+	took := time.Since(startTime).Milliseconds()
 
 	stdout = truncateOutput(stdout)
 	stderr = truncateOutput(stderr)
@@ -304,10 +310,13 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 		stdout += "\n" + errorMessage
 	}
 
+	metadata := BashToolResponseMetadata{
+		Took: took,
+	}
 	if stdout == "" {
-		return NewTextResponse("no output"), nil
+		return WithResponseMetadata(NewTextResponse("no output"), metadata), nil
 	}
-	return NewTextResponse(stdout), nil
+	return WithResponseMetadata(NewTextResponse(stdout), metadata), nil
 }
 
 func truncateOutput(content string) string {
diff --git a/internal/llm/tools/tools.go b/internal/llm/tools/tools.go
index e15c1c31f..6bb528686 100644
--- a/internal/llm/tools/tools.go
+++ b/internal/llm/tools/tools.go
@@ -1,6 +1,9 @@
 package tools
 
-import "context"
+import (
+	"context"
+	"encoding/json"
+)
 
 type ToolInfo struct {
 	Name        string
@@ -17,9 +20,10 @@ const (
 )
 
 type ToolResponse struct {
-	Type    toolResponseType `json:"type"`
-	Content string           `json:"content"`
-	IsError bool             `json:"is_error"`
+	Type     toolResponseType `json:"type"`
+	Content  string           `json:"content"`
+	Metadata string           `json:"metadata,omitempty"`
+	IsError  bool             `json:"is_error"`
 }
 
 func NewTextResponse(content string) ToolResponse {
@@ -29,6 +33,17 @@ func NewTextResponse(content string) ToolResponse {
 	}
 }
 
+func WithResponseMetadata(response ToolResponse, metadata any) ToolResponse {
+	if metadata != nil {
+		metadataBytes, err := json.Marshal(metadata)
+		if err != nil {
+			return response
+		}
+		response.Metadata = string(metadataBytes)
+	}
+	return response
+}
+
 func NewTextErrorResponse(content string) ToolResponse {
 	return ToolResponse{
 		Type:    ToolResponseTypeText,