From 89343d2e05cfe17f1b5903db2dce2e48d7749cc4 Mon Sep 17 00:00:00 2001
From: Grant Linville <grant@acorn.io>
Date: Fri, 14 Jun 2024 11:08:28 -0400
Subject: [PATCH] exclude output of cred tools in CallFinish events

Signed-off-by: Grant Linville <grant@acorn.io>
---
 pkg/runner/monitor.go | 19 -------------------
 pkg/runner/runner.go  | 12 ++++++++++--
 2 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/pkg/runner/monitor.go b/pkg/runner/monitor.go
index 0efa79b5..48e64b0b 100644
--- a/pkg/runner/monitor.go
+++ b/pkg/runner/monitor.go
@@ -28,22 +28,3 @@ func (n noopMonitor) Stop(string, error) {}
 func (n noopMonitor) Pause() func() {
 	return func() {}
 }
-
-type credWrapper struct {
-	m Monitor
-}
-
-func (c credWrapper) Event(e Event) {
-	if e.Type == EventTypeCallFinish {
-		e.Content = "credential tool output redacted"
-	}
-	c.m.Event(e)
-}
-
-func (c credWrapper) Stop(s string, err error) {
-	c.m.Stop(s, err)
-}
-
-func (c credWrapper) Pause() func() {
-	return c.m.Pause()
-}
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
index d52236e2..7f04408b 100644
--- a/pkg/runner/runner.go
+++ b/pkg/runner/runner.go
@@ -542,7 +542,7 @@ func (r *Runner) resume(callCtx engine.Context, monitor Monitor, env []string, s
 				Time:        time.Now(),
 				CallContext: callCtx.GetCallContext(),
 				Type:        EventTypeCallFinish,
-				Content:     *state.Continuation.Result,
+				Content:     getFinishEventContent(*state, callCtx),
 			})
 			if callCtx.Tool.Chat {
 				return &State{
@@ -786,6 +786,14 @@ func (r *Runner) subCalls(callCtx engine.Context, monitor Monitor, env []string,
 	return state, callResults, nil
 }
 
+func getFinishEventContent(state State, callCtx engine.Context) string {
+	if callCtx.ToolCategory == engine.CredentialToolCategory {
+		return ""
+	}
+
+	return *state.Continuation.Result
+}
+
 func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env []string) ([]string, error) {
 	// Since credential tools (usually) prompt the user, we want to only run one at a time.
 	r.credMutex.Lock()
@@ -858,7 +866,7 @@ func (r *Runner) handleCredentials(callCtx engine.Context, monitor Monitor, env
 				return nil, fmt.Errorf("failed to create subcall context for tool %s: %w", credToolName, err)
 			}
 
-			res, err := r.call(subCtx, credWrapper{monitor}, env, input)
+			res, err := r.call(subCtx, monitor, env, input)
 			if err != nil {
 				return nil, fmt.Errorf("failed to run credential tool %s: %w", credToolName, err)
 			}