diff --git a/cmd/root/api.go b/cmd/root/api.go index 9498ed821..0ed1f5c06 100644 --- a/cmd/root/api.go +++ b/cmd/root/api.go @@ -146,6 +146,11 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error { if err != nil { return fmt.Errorf("creating session store: %w", err) } + defer func() { + if err := sessionStore.Close(); err != nil { + slog.Error("Failed to close session store", "error", err) + } + }() sources, err := config.ResolveSources(agentsPath, f.runConfig.EnvProvider()) if err != nil { diff --git a/cmd/root/run.go b/cmd/root/run.go index f959a06d2..a59c0bd01 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -240,6 +240,9 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s if err := loadResult.Team.StopToolSets(cleanupCtx); err != nil { slog.Error("Failed to stop tool sets", "error", err) } + if err := rt.Close(); err != nil { + slog.Error("Failed to close runtime", "error", err) + } } } defer cleanup() diff --git a/pkg/app/app_test.go b/pkg/app/app_test.go index f2ec8a89f..3977d9529 100644 --- a/pkg/app/app_test.go +++ b/pkg/app/app_test.go @@ -59,6 +59,7 @@ func (m *mockRuntime) UpdateSessionTitle(_ context.Context, sess *session.Sessio return nil } func (m *mockRuntime) TitleGenerator() *sessiontitle.Generator { return nil } +func (m *mockRuntime) Close() error { return nil } func (m *mockRuntime) Stop() {} // Verify mockRuntime implements runtime.Runtime diff --git a/pkg/runtime/commands_test.go b/pkg/runtime/commands_test.go index 55d5f1b22..1221aec26 100644 --- a/pkg/runtime/commands_test.go +++ b/pkg/runtime/commands_test.go @@ -64,6 +64,7 @@ func (m *mockRuntime) UpdateSessionTitle(context.Context, *session.Session, stri return nil } func (m *mockRuntime) TitleGenerator() *sessiontitle.Generator { return nil } +func (m *mockRuntime) Close() error { return nil } func (m *mockRuntime) RegenerateTitle(context.Context, *session.Session, chan Event) { } diff --git a/pkg/runtime/remote_runtime.go b/pkg/runtime/remote_runtime.go index 049091169..4cb7f0960 100644 --- a/pkg/runtime/remote_runtime.go +++ b/pkg/runtime/remote_runtime.go @@ -452,4 +452,9 @@ func (r *RemoteRuntime) TitleGenerator() *sessiontitle.Generator { return nil } +// Close is a no-op for remote runtimes. +func (r *RemoteRuntime) Close() error { + return nil +} + var _ Runtime = (*RemoteRuntime)(nil) diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index e8610106a..ac8e6c92b 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -153,6 +153,9 @@ type Runtime interface { // TitleGenerator returns a generator for automatic session titles, or nil // if the runtime does not support local title generation (e.g. remote runtimes). TitleGenerator() *sessiontitle.Generator + + // Close releases resources held by the runtime (e.g., session store connections). + Close() error } // PermissionsInfo contains the allow and deny patterns for tool permissions. @@ -671,6 +674,14 @@ func (r *LocalRuntime) SessionStore() session.Store { return r.sessionStore } +// Close releases resources held by the runtime, including the session store. +func (r *LocalRuntime) Close() error { + if r.sessionStore != nil { + return r.sessionStore.Close() + } + return nil +} + // UpdateSessionTitle persists the session title via the session store. func (r *LocalRuntime) UpdateSessionTitle(ctx context.Context, sess *session.Session, title string) error { sess.Title = title diff --git a/pkg/session/store.go b/pkg/session/store.go index ccfd2d305..ec58ac94c 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -108,6 +108,9 @@ type Store interface { // UpdateSessionTitle updates only the title UpdateSessionTitle(ctx context.Context, sessionID, title string) error + + // Close releases any resources held by the store (e.g., database connections). + Close() error } type InMemorySessionStore struct { @@ -350,6 +353,11 @@ func (s *InMemorySessionStore) UpdateSessionTitle(_ context.Context, sessionID, return nil } +// Close is a no-op for in-memory stores. +func (s *InMemorySessionStore) Close() error { + return nil +} + // NewSQLiteSessionStore creates a new SQLite session store func NewSQLiteSessionStore(path string) (Store, error) { store, err := openAndMigrateSQLiteStore(path)