odidere/internal/job/job_test.go

package job

import (
	"database/sql"
	"encoding/json/v2"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"code.chimeric.al/chimerical/odidere/internal/config"
	"code.chimeric.al/chimerical/odidere/internal/db"
	"code.chimeric.al/chimerical/odidere/internal/llm"
	"code.chimeric.al/chimerical/odidere/internal/tool"
)

func testDependencies(t *testing.T) (*db.SQLite, *llm.Client, *tool.Registry) {
	t.Helper()

	database, err := db.Connect(t.Context(), db.Config{
		Path: filepath.Join(t.TempDir(), "job.db"),
	})
	if err != nil {
		t.Fatalf("db.Connect: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	if err := database.ApplySchema(t.Context(), slog.Default()); err != nil {
		t.Fatalf("ApplySchema: %v", err)
	}
	client, err := llm.NewClient(
		llm.Config{URL: "http://localhost"}, slog.Default(),
	)
	if err != nil {
		t.Fatalf("llm.NewClient: %v", err)
	}
	registry, err := tool.NewRegistry(nil)
	if err != nil {
		t.Fatalf("tool.NewRegistry: %v", err)
	}
	return database, client, registry
}

func TestConfigValidate(t *testing.T) {
	tests := []struct {
		name    string
		config  config.JobConfig
		wantErr bool
	}{
		{
			name:    "empty config",
			config:  config.JobConfig{},
			wantErr: true,
		},
		{
			name: "missing name",
			config: config.JobConfig{
				Schedule: "0 9 * * *",
				Task:     "do something",
				UserId:   "default",
			},
			wantErr: true,
		},
		{
			name: "missing task",
			config: config.JobConfig{
				Name:     "test",
				Schedule: "0 9 * * *",
				UserId:   "default",
			},
			wantErr: true,
		},
		{
			name: "missing schedule",
			config: config.JobConfig{
				Name:   "test",
				Task:   "do something",
				UserId: "default",
			},
			wantErr: true,
		},
		{
			name: "missing user ID",
			config: config.JobConfig{
				Name:     "test",
				Task:     "do something",
				Schedule: "0 9 * * *",
			},
			wantErr: true,
		},
		{
			name: "negative max iterations",
			config: config.JobConfig{
				Name:          "test",
				Task:          "do something",
				Schedule:      "0 9 * * *",
				MaxIterations: -1,
				UserId:        "default",
			},
			wantErr: true,
		},
		{
			name: "minimal valid config",
			config: config.JobConfig{
				Name:     "test",
				Task:     "do something",
				Schedule: "0 9 * * *",
				UserId:   "default",
			},
			wantErr: false,
		},
		{
			name: "full valid config",
			config: config.JobConfig{
				Name:          "test",
				Task:          "do something",
				Schedule:      "0 9 * * *",
				Timeout:       10 * time.Minute,
				MaxIterations: 20,
				SystemMessage: "You are an assistant.",
				UserId:        "default",
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := tt.config.Validate()
			if (err != nil) != tt.wantErr {
				t.Errorf("Validate() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestConfigMaxIterations(t *testing.T) {
	// Zero means unlimited
	cfg := config.JobConfig{Name: "test"}
	if got := cfg.MaxIterations; got != 0 {
		t.Errorf("default max iterations = %d, want 0", got)
	}

	cfg.MaxIterations = 0
	if got := cfg.MaxIterations; got != 0 {
		t.Errorf("explicit zero max iterations = %d, want 0", got)
	}

	// Positive value passes through
	cfg.MaxIterations = 25
	if got := cfg.MaxIterations; got != 25 {
		t.Errorf("custom max iterations = %d, want 25", got)
	}
}

func TestConfigTimeout(t *testing.T) {
	// A zero Timeout means the job has no per-job execution limit;
	// Run then relies on the callers context.
	var cfg config.JobConfig
	if cfg.Timeout != 0 {
		t.Errorf("default timeout = %v, want 0", cfg.Timeout)
	}
}

func TestNew(t *testing.T) {
	database, client, registry := testDependencies(t)
	cfg := config.JobConfig{
		Name:          "test",
		Task:          "do something",
		Schedule:      "0 9 * * *",
		Timeout:       10 * time.Minute,
		MaxIterations: 20,
		SystemMessage: "You are an assistant.",
	}

	cfg.Model.Model = "test-model"
	cfg.UserId = "default"
	j, err := New(cfg, database.Writer, client, slog.Default(), registry)
	if err != nil {
		t.Fatalf("New() error = %v", err)
	}

	if j.Name() != "test" {
		t.Errorf("Name() = %q, want %q", j.Name(), "test")
	}
	if j.Schedule() != "0 9 * * *" {
		t.Errorf("Schedule() = %q, want %q", j.Schedule(), "0 9 * * *")
	}
	if j.SystemMessage() != "You are an assistant." {
		t.Errorf("SystemMessage() = %q, want %q", j.SystemMessage(), "You are an assistant.")
	}
	if j.Task() != "do something" {
		t.Errorf("Task() = %q, want %q", j.Task(), "do something")
	}
	if j.Timeout().Minutes() != 10 {
		t.Errorf("Timeout() = %v, want 10m", j.Timeout())
	}
	if j.MaxIterations() != 20 {
		t.Errorf("MaxIterations() = %d, want 20", j.MaxIterations())
	}
	if j.Model() != "test-model" {
		t.Errorf("Model() = %q, want %q", j.Model(), "test-model")
	}
	if j.Cron() == nil {
		t.Error("Cron() = nil, want non-nil")
	}
}

func TestNew_InvalidSchedule(t *testing.T) {
	database, client, registry := testDependencies(t)
	cfg := config.JobConfig{
		Name:     "test",
		Task:     "do something",
		Schedule: "not-a-cron",
	}

	cfg.Model.Model = "test-model"
	cfg.UserId = "default"
	_, err := New(cfg, database.Writer, client, slog.Default(), registry)
	if err == nil {
		t.Fatal("New() expected error for invalid schedule")
	}
}

func TestNew_InvalidTimeout(t *testing.T) {
	database, client, registry := testDependencies(t)
	cfg := config.JobConfig{
		Name:     "test",
		Task:     "do something",
		Schedule: "0 9 * * *",
		Timeout:  -time.Minute,
	}

	cfg.Model.Model = "test-model"
	cfg.UserId = "default"
	_, err := New(cfg, database.Writer, client, slog.Default(), registry)
	if err == nil {
		t.Fatal("New() expected error for invalid timeout")
	}
}

func TestNewRequiresDependencies(t *testing.T) {
	database, client, registry := testDependencies(t)
	cfg := config.JobConfig{
		Name:     "test",
		Task:     "do something",
		Schedule: "0 9 * * *",
	}
	for _, test := range []struct {
		name     string
		registry *tool.Registry
		database *sql.DB
		client   *llm.Client
		logger   *slog.Logger
		userID   string
	}{
		{
			name:     "nil registry",
			database: database.Writer,
			client:   client,
			logger:   slog.Default(),
			userID:   "default",
		},
		{
			name:     "nil database",
			registry: registry,
			client:   client,
			logger:   slog.Default(),
			userID:   "default",
		},
		{
			name:     "empty user ID",
			registry: registry,
			database: database.Writer,
			client:   client,
			logger:   slog.Default(),
		},
		{
			name:     "nil llm client",
			registry: registry,
			database: database.Writer,
			logger:   slog.Default(),
			userID:   "default",
		},
		{
			name:     "nil logger",
			registry: registry,
			database: database.Writer,
			client:   client,
			userID:   "default",
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			cfg.UserId = test.userID
			_, err := New(
				cfg, test.database, test.client, test.logger,
				test.registry,
			)
			if err == nil {
				t.Fatal("New() expected error")
			}
		})
	}
}

func TestRunPersistsToolLoop(t *testing.T) {
	var requests atomic.Int32
	inference := httptest.NewServer(http.HandlerFunc(func(
		w http.ResponseWriter, r *http.Request,
	) {
		var request llm.ChatRequest
		if err := json.UnmarshalRead(r.Body, &request); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		switch requests.Add(1) {
		case 1:
			if len(request.Messages) != 2 ||
				request.Messages[0].Role != llm.RoleSystem ||
				request.Messages[1].Role != llm.RoleUser {
				http.Error(w, "unexpected initial history", http.StatusBadRequest)
				return
			}
			fmt.Fprint(w, `{
				"choices":[{
					"message":{
						"role":"assistant",
						"tool_calls":[{
							"id":"call_test",
							"type":"function",
							"function":{"name":"check_db","arguments":"{}"}
						}]
				},
				"finish_reason":"tool_calls"
				}]
			}`)
		case 2:
			if len(request.Messages) != 4 ||
				request.Messages[2].Role != llm.RoleAssistant ||
				request.Messages[3].Role != llm.RoleTool ||
				request.Messages[3].ToolCallID != "call_test" {
				http.Error(w, "unexpected tool history", http.StatusBadRequest)
				return
			}
			fmt.Fprint(w, `{
				"choices":[{
					"message":{"role":"assistant","content":"finished"},
					"finish_reason":"stop"
				}]
			}`)
		default:
			http.Error(w, "unexpected inference request", http.StatusBadRequest)
		}
	}))
	defer inference.Close()

	var (
		ctx          = t.Context()
		log          = slog.New(slog.NewTextHandler(io.Discard, nil))
		databasePath = filepath.Join(t.TempDir(), "job.db")
	)
	database, err := db.Connect(ctx, db.Config{Path: databasePath})
	if err != nil {
		t.Fatalf("db.Connect: %v", err)
	}
	defer database.Close()
	if err := database.ApplySchema(ctx, slog.Default()); err != nil {
		t.Fatalf("ApplySchema: %v", err)
	}
	t.Setenv("ODIDERE_JOB_TEST_DATABASE", databasePath)
	executable, err := os.Executable()
	if err != nil {
		t.Fatalf("os.Executable: %v", err)
	}
	registry, err := tool.NewRegistry([]config.ToolConfig{{
		Name:        "check_db",
		Description: "Verify job persistence",
		Command:     executable,
		Arguments:   []string{"-test.run=^TestJobPersistenceHelper$"},
		Parameters:  map[string]any{"type": "object"},
	}})
	if err != nil {
		t.Fatalf("tool.NewRegistry: %v", err)
	}
	client, err := llm.NewClient(llm.Config{
		Provider: "test",
		URL:      inference.URL,
	}, log)
	if err != nil {
		t.Fatalf("llm.NewClient: %v", err)
	}
	job, err := New(
		config.JobConfig{
			Name:          "daily report",
			Task:          "prepare the report",
			Schedule:      "0 9 * * *",
			SystemMessage: "system",
			Model:         config.ModelConfig{Model: "test-model"},
			UserId:        "default",
		},
		database.Writer,
		client,
		log,
		registry,
	)
	if err != nil {
		t.Fatalf("New: %v", err)
	}

	result := job.Run(ctx)
	if result.Error != nil {
		t.Fatalf("Run: %v", result.Error)
	}
	if result.ConversationID == "" {
		t.Fatal("conversation ID is empty")
	}
	if requests.Load() != 2 {
		t.Fatalf("inference requests = %d, want 2", requests.Load())
	}

	q := db.New(database.Reader)
	conversation, err := q.SelectConversationId(ctx, result.ConversationID)
	if err != nil {
		t.Fatalf("SelectConversationId: %v", err)
	}
	if conversation.UserId != "default" {
		t.Errorf("conversation user = %q, want default", conversation.UserId)
	}
	if conversation.Title == nil || *conversation.Title != "daily report" {
		t.Errorf("conversation title = %v, want daily report", conversation.Title)
	}
	stored, err := q.SelectMessagesConversationId(ctx, result.ConversationID)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	wantRoles := []db.MessageRole{
		db.MessageRoleSystem,
		db.MessageRoleUser,
		db.MessageRoleAssistant,
		db.MessageRoleTool,
		db.MessageRoleAssistant,
	}
	if len(stored) != len(wantRoles) {
		t.Fatalf("stored messages = %d, want %d", len(stored), len(wantRoles))
	}
	for i, message := range stored {
		if message.Role != db.MessageRoleAssistant {
			if message.Metadata != nil {
				t.Errorf("message %d metadata = %v, want nil", i, message.Metadata)
			}
			continue
		}
		if message.Metadata == nil {
			t.Fatalf("assistant %d has no metadata", i)
		}
		var metadata llm.Metadata
		if err := json.Unmarshal(message.Metadata, &metadata); err != nil {
			t.Fatal(err)
		}
		if metadata.Provider == nil || *metadata.Provider != "test" ||
			metadata.RequestedModel == nil || *metadata.RequestedModel != "test-model" {
			t.Errorf("assistant %d metadata = %s", i, string(message.Metadata))
		}
		wantReason := "stop"
		if i == 2 {
			wantReason = "tool_calls"
		}
		if metadata.FinishReason == nil || *metadata.FinishReason != wantReason {
			t.Errorf("assistant %d finish reason = %v", i, metadata.FinishReason)
		}
	}
	for i, want := range wantRoles {
		if stored[i].Sequence != i || stored[i].Role != want {
			t.Errorf(
				"stored message %d = sequence %d, role %q; want %d, %q",
				i, stored[i].Sequence, stored[i].Role, i, want,
			)
		}
	}
}

func TestJobPersistenceHelper(t *testing.T) {
	path := os.Getenv("ODIDERE_JOB_TEST_DATABASE")
	if path == "" {
		return
	}
	database, err := sql.Open("sqlite3", "file:"+path+"?mode=ro")
	if err != nil {
		t.Fatalf("sql.Open: %v", err)
	}
	defer database.Close()

	var role string
	if err := database.QueryRow(
		"SELECT role FROM messages ORDER BY created_at DESC, sequence DESC LIMIT 1",
	).Scan(&role); err != nil {
		t.Fatalf("select latest message: %v", err)
	}
	if role != string(db.MessageRoleAssistant) {
		t.Fatalf("latest message role = %q, want assistant", role)
	}
}

func TestCallTool(t *testing.T) {
	database, client, _ := testDependencies(t)
	registry, err := tool.NewRegistry([]config.ToolConfig{
		{
			Name:        "finish",
			Description: "finish",
			Command:     "sh",
			Arguments:   []string{"-c", "echo done"},
		},
		{
			Name:        "fail",
			Description: "fail",
			Command:     "sh",
			Arguments:   []string{"-c", "echo oops; exit 3"},
		},
		{
			Name:        "missing",
			Description: "missing",
			Command:     "nonexistent_binary_xyz",
		},
		{
			Name:        "interrupt",
			Description: "interrupt",
			Command:     "sh",
			// exec replaces the shell, so the timeout kills the
			// process that holds the output pipe.
			Arguments: []string{"-c", "echo partial; exec sleep 10"},
			Timeout:   100 * time.Millisecond,
		},
	})
	if err != nil {
		t.Fatalf("tool.NewRegistry: %v", err)
	}
	log := slog.New(slog.DiscardHandler)
	job, err := New(
		config.JobConfig{
			Name:     "call tool",
			Task:     "call a tool",
			Schedule: "0 9 * * *",
			UserId:   "default",
		},
		database.Writer,
		client,
		log,
		registry,
	)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	conversation, err := job.createConversation(t.Context(), nil)
	if err != nil {
		t.Fatalf("createConversation: %v", err)
	}

	tests := []struct {
		name string
		tool string
		// want lists text the message must contain, in this order.
		want []string
		// warn reports whether the message must carry
		// interruptedWarning.
		warn bool
		// status is the execution row's status
		// after callTool returns.
		status db.ExecutionStatus
	}{
		{
			name:   "finished",
			tool:   "finish",
			want:   []string{"▸ stdout:\ndone\n", "┗ exit: 0"},
			status: db.ExecutionStatusCompleted,
		},
		{
			name:   "non-zero exit",
			tool:   "fail",
			want:   []string{"▸ stdout:\noops\n", "┗ exit: 3"},
			status: db.ExecutionStatusCompleted,
		},
		{
			name:   "never started",
			tool:   "missing",
			want:   []string{"executable file not found"},
			status: db.ExecutionStatusCompleted,
		},
		{
			name: "interrupted",
			tool: "interrupt",
			want: []string{
				"timed out",
				interruptedWarning,
				"▸ stdout:\npartial\n",
			},
			warn:   true,
			status: db.ExecutionStatusRunning,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			call := llm.ToolCall{
				Function: llm.Function{Arguments: "{}", Name: tt.tool},
				ID:       "call_" + tt.tool,
				Type:     llm.ToolTypeFunction,
			}
			assistant, err := job.insertMessage(
				t.Context(), conversation.Id, llm.Message{
					Role:      llm.RoleAssistant,
					ToolCalls: []llm.ToolCall{call},
				},
				&llm.Metadata{ID: new("assistant-completion")},
				nil,
			)
			if err != nil {
				t.Fatalf("insertMessage: %v", err)
			}

			message, err := job.callTool(
				t.Context(), log, conversation.Id, assistant.Id, call,
			)
			if err != nil {
				t.Fatalf("callTool: %v", err)
			}
			if message.Role != llm.RoleTool ||
				message.ToolCallID != call.ID ||
				message.Name != tt.tool {
				t.Errorf(
					"message = role %q, tool call %q, name %q; "+
						"want %q, %q, %q",
					message.Role, message.ToolCallID, message.Name,
					llm.RoleTool, call.ID, tt.tool,
				)
			}
			content := message.ContentParts[0].Text
			rest := content
			for _, want := range tt.want {
				i := strings.Index(rest, want)
				if i < 0 {
					t.Fatalf(
						"content missing %q, or out of order:\n%s",
						want, content,
					)
				}
				rest = rest[i+len(want):]
			}
			if got := strings.Contains(
				content, interruptedWarning,
			); got != tt.warn {
				t.Errorf(
					"content has warning = %t, want %t:\n%s",
					got, tt.warn, content,
				)
			}

			stored, err := db.New(database.Reader).
				SelectMessagesConversationId(
					t.Context(), conversation.Id,
				)
			if err != nil {
				t.Fatalf("SelectMessagesConversationId: %v", err)
			}
			want, err := json.Marshal(*message)
			if err != nil {
				t.Fatalf("json.Marshal: %v", err)
			}
			wantMetadata := `{"choice_index":null,"created":null,` +
				`"finish_reason":null,"id":"assistant-completion",` +
				`"logprobs":null,"model":null,"provider":null,` +
				`"requested_model":null,"service_tier":null,` +
				`"system_fingerprint":null,"timings":null,"usage":null}`
			if assistant.Metadata == nil ||
				string(assistant.Metadata) != wantMetadata {
				t.Errorf("assistant metadata = %v", assistant.Metadata)
			}
			last := stored[len(stored)-1]
			if last.Metadata != nil {
				t.Errorf("tool metadata = %v, want nil", last.Metadata)
			}
			if last.Role != db.MessageRoleTool ||
				last.Content != string(want) {
				t.Errorf(
					"last stored message = role %q, %s; "+
						"want %q, %s",
					last.Role, last.Content,
					db.MessageRoleTool, want,
				)
			}

			executions, err := db.New(database.Reader).
				SelectExecutionsMessageId(t.Context(), assistant.Id)
			if err != nil {
				t.Fatalf("SelectExecutionsMessageId: %v", err)
			}
			if len(executions) != 1 ||
				executions[0].ToolCallId != call.ID ||
				executions[0].Status != tt.status {
				t.Errorf(
					"executions = %+v, want one %q row for %q",
					executions, tt.status, call.ID,
				)
			}
		})
	}
}

// TestCallToolUnsavedRow checks that callTool does not run the tool when
// its executions row cannot be saved. The row references the assistant
// message, so a message ID that does not exist makes the save fail.
func TestCallToolUnsavedRow(t *testing.T) {
	database, client, _ := testDependencies(t)
	marker := filepath.Join(t.TempDir(), "ran")
	t.Setenv("ODIDERE_CALL_TOOL_MARKER", marker)
	registry, err := tool.NewRegistry([]config.ToolConfig{{
		Name:        "touch",
		Description: "touch",
		Command:     "sh",
		Arguments:   []string{"-c", `touch "$ODIDERE_CALL_TOOL_MARKER"`},
	}})
	if err != nil {
		t.Fatalf("tool.NewRegistry: %v", err)
	}
	log := slog.New(slog.DiscardHandler)
	job, err := New(
		config.JobConfig{
			Name:     "unsaved row",
			Task:     "call a tool",
			Schedule: "0 9 * * *",
			UserId:   "default",
		},
		database.Writer,
		client,
		log,
		registry,
	)
	if err != nil {
		t.Fatalf("New: %v", err)
	}

	conversation, err := job.createConversation(t.Context(), nil)
	if err != nil {
		t.Fatalf("createConversation: %v", err)
	}

	message, err := job.callTool(
		t.Context(), log, conversation.Id, "missing",
		llm.ToolCall{
			Function: llm.Function{
				Arguments: "{}", Name: "touch",
			},
			ID:   "call_touch",
			Type: llm.ToolTypeFunction,
		},
	)
	if err == nil {
		t.Fatal("callTool: expected error")
	}
	if message != nil {
		t.Errorf("message = %+v, want nil", message)
	}
	if _, err := os.Stat(marker); err == nil {
		t.Error("tool ran without a saved row")
	}
	executions, err := db.New(database.Reader).
		SelectExecutionsMessageId(t.Context(), "missing")
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId: %v", err)
	}
	if len(executions) != 0 {
		t.Errorf("executions = %+v, want none", executions)
	}
}